pyobvector 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

pyobvector/client/milvus_like_client.py

@@ -1,4 +1,5 @@
 """Milvus Like Client."""
+
 import logging
 import json
 from typing import Optional, Union
@@ -62,11 +63,11 @@ class MilvusLikeClient(Client):
         index_params: Optional[IndexParams] = None,  # Used for custom setup
         max_length: int = 16384,
         **kwargs,
-    ): # pylint: disable=unused-argument
-        """Create a collection. 
+    ):  # pylint: disable=unused-argument
+        """Create a collection.
         If `schema` is not None, `dimension`, `primary_field_name`, `id_type`, `vector_field_name`,
         `metric_type`, `auto_id` will be ignored.
-        
+
         Args:
             collection_name (string): collection name
             dimension (Optional[int]): vector data dimension
@@ -146,10 +147,12 @@ class MilvusLikeClient(Client):
             )
 
     def get_collection_stats(
-        self, collection_name: str, timeout: Optional[float] = None # pylint: disable=unused-argument
+        self,
+        collection_name: str,
+        timeout: Optional[float] = None,  # pylint: disable=unused-argument
     ) -> dict:
         """Get collection row count.
-        
+
         Args:
             collection_name (string): collection name
             timeout (Optional[float]): not used in OceanBase
@@ -166,8 +169,10 @@ class MilvusLikeClient(Client):
                 return {"row_count": cnt}
 
     def has_collection(
-        self, collection_name: str, timeout: Optional[float] = None # pylint: disable=unused-argument
-    ) -> bool: # pylint: disable=unused-argument
+        self,
+        collection_name: str,
+        timeout: Optional[float] = None,  # pylint: disable=unused-argument
+    ) -> bool:  # pylint: disable=unused-argument
         """Check if collection exists.
 
         Args:
@@ -181,17 +186,20 @@ class MilvusLikeClient(Client):
 
     def drop_collection(self, collection_name: str) -> None:
         """Drop collection if exists.
-        
+
         Args:
             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
+        self,
+        old_name: str,
+        new_name: str,
+        timeout: Optional[float] = None,  # pylint: disable=unused-argument
     ) -> None:
         """Rename collection.
-        
+
         Args:
             old_name (string): old collection name
             new_name (string): new collection name
@@ -206,7 +214,7 @@ class MilvusLikeClient(Client):
         collection_name: str,
     ):
         """Load table into SQLAlchemy metadata.
-        
+
         Args:
             collection_name (string): which collection to load
 
@@ -230,9 +238,9 @@ class MilvusLikeClient(Client):
         index_params: IndexParams,
         timeout: Optional[float] = None,
         **kwargs,
-    ): # pylint: disable=unused-argument
+    ):  # pylint: disable=unused-argument
         """Create vector index with index params.
-        
+
         Args:
             collection_name (string): which collection to create vector index
             index_params (IndexParams): the vector index parameters
@@ -263,9 +271,9 @@ class MilvusLikeClient(Client):
         index_name: str,
         timeout: Optional[float] = None,
         **kwargs,
-    ): # pylint: disable=unused-argument
+    ):  # pylint: disable=unused-argument
         """Drop index on specified collection.
-        
+
         If the index not exists, SQL ERROR 1091 will raise.
 
         Args:
@@ -283,7 +291,7 @@ class MilvusLikeClient(Client):
         trigger_threshold: int = 10000,
     ):
         """Refresh vector index for performance.
-        
+
         Args:
             collection_name (string): collection name
             index_name (string): vector index name
@@ -303,7 +311,7 @@ class MilvusLikeClient(Client):
         trigger_threshold: float = 0.2,
     ):
         """Rebuild vector index for performance.
-        
+
         Args:
             collection_name (string): collection name
             index_name (string): vector index name
@@ -356,13 +364,13 @@ class MilvusLikeClient(Client):
         limit: int = 10,
         output_fields: Optional[list[str]] = None,
         search_params: Optional[dict] = None,
-        timeout: Optional[float] = None, # pylint: disable=unused-argument
+        timeout: Optional[float] = None,  # pylint: disable=unused-argument
         partition_names: Optional[list[str]] = None,
-        **kwargs, # pylint: disable=unused-argument
+        **kwargs,  # pylint: disable=unused-argument
     ) -> list[dict]:
         """Perform ann search.
         Note: OceanBase does not support batch search now. `data` & the return value is not a batch.
-        
+
         Args:
             collection_name (string): collection name
             data (list): the vector/sparse_vector data to search
@@ -392,9 +400,7 @@ class MilvusLikeClient(Client):
                         message=ExceptionsMessage.MetricTypeParamTypeInvalid,
                     )
                 lower_metric_type_str = search_params["metric_type"].lower()
-                if lower_metric_type_str not in (
-                    "l2", "neg_ip", "cosine", "ip"
-                ):
+                if lower_metric_type_str not in ("l2", "neg_ip", "cosine", "ip"):
                     raise VectorMetricTypeException(
                         code=ErrorCode.INVALID_ARGUMENT,
                         message=ExceptionsMessage.MetricTypeValueInvalid,
@@ -416,25 +422,34 @@ class MilvusLikeClient(Client):
 
         if with_dist:
             if isinstance(data, list):
-                columns.append(distance_func(table.c[anns_field],
-                    "[" + ",".join([str(np.float32(v)) for v in data]) + "]"))
+                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)
-        
+
         if isinstance(data, list):
-            stmt = stmt.order_by(distance_func(table.c[anns_field],
-                "[" + ",".join([str(np.float32(v)) for v in data]) + "]"))
+            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,
-                compile_kwargs={"literal_binds": True}
-            ))
+            str(
+                stmt.compile(
+                    dialect=self.engine.dialect, compile_kwargs={"literal_binds": True}
+                )
+            )
             + f" APPROXIMATE limit {limit}"
         )
 
@@ -468,12 +483,12 @@ class MilvusLikeClient(Client):
         collection_name: str,
         flter=None,
         output_fields: Optional[list[str]] = None,
-        timeout: Optional[float] = None, # pylint: disable=unused-argument
+        timeout: Optional[float] = None,  # pylint: disable=unused-argument
         partition_names: Optional[list[str]] = None,
-        **kwargs, # pylint: disable=unused-argument
+        **kwargs,  # pylint: disable=unused-argument
     ) -> list[dict]:
         """Query records.
-        
+
         Args:
             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)
@@ -508,10 +523,12 @@ class MilvusLikeClient(Client):
                 if partition_names is None:
                     execute_res = conn.execute(stmt)
                 else:
-                    stmt_str = str(stmt.compile(
-                        dialect=self.engine.dialect,
-                        compile_kwargs={"literal_binds": True}
-                    ))
+                    stmt_str = str(
+                        stmt.compile(
+                            dialect=self.engine.dialect,
+                            compile_kwargs={"literal_binds": True},
+                        )
+                    )
                     stmt_str = self._insert_partition_hint_for_query_sql(
                         stmt_str, f"PARTITION({', '.join(partition_names)})"
                     )
@@ -534,12 +551,12 @@ class MilvusLikeClient(Client):
         collection_name: str,
         ids: Union[list, str, int] = None,
         output_fields: Optional[list[str]] = None,
-        timeout: Optional[float] = None, # pylint: disable=unused-argument
+        timeout: Optional[float] = None,  # pylint: disable=unused-argument
         partition_names: Optional[list[str]] = None,
-        **kwargs, # pylint: disable=unused-argument
+        **kwargs,  # pylint: disable=unused-argument
     ) -> list[dict]:
         """Get records with specified primary field `ids`.
-        
+
         Args:
             collection_name (string): collection name
             ids (Union[list, str, int]): specified primary field values
@@ -586,10 +603,12 @@ class MilvusLikeClient(Client):
                 if partition_names is None:
                     execute_res = conn.execute(stmt)
                 else:
-                    stmt_str = str(stmt.compile(
-                        dialect=self.engine.dialect,
-                        compile_kwargs={"literal_binds": True}
-                    ))
+                    stmt_str = str(
+                        stmt.compile(
+                            dialect=self.engine.dialect,
+                            compile_kwargs={"literal_binds": True},
+                        )
+                    )
                     stmt_str = self._insert_partition_hint_for_query_sql(
                         stmt_str, f"PARTITION({', '.join(partition_names)})"
                     )
@@ -611,10 +630,10 @@ class MilvusLikeClient(Client):
         self,
         collection_name: str,
         ids: Optional[Union[list, str, int]] = None,
-        timeout: Optional[float] = None, # pylint: disable=unused-argument
+        timeout: Optional[float] = None,  # pylint: disable=unused-argument
         flter=None,
         partition_name: Optional[str] = "",
-        **kwargs, # pylint: disable=unused-argument
+        **kwargs,  # pylint: disable=unused-argument
     ) -> dict:
         """Delete data in collection.
 
@@ -675,11 +694,9 @@ class MilvusLikeClient(Client):
         data: Union[dict, list[dict]],
         timeout: Optional[float] = None,
         partition_name: Optional[str] = "",
-    ) -> (
-        None
-    ):  # pylint: disable=unused-argument
+    ) -> None:  # pylint: disable=unused-argument
         """Insert data into collection.
-        
+
         Args:
             collection_name (string): collection name
             data (Union[Dict, List[Dict]]): data that will be inserted
@@ -701,11 +718,11 @@ class MilvusLikeClient(Client):
         self,
         collection_name: str,
         data: Union[dict, list[dict]],
-        timeout: Optional[float] = None, # pylint: disable=unused-argument
+        timeout: Optional[float] = None,  # pylint: disable=unused-argument
         partition_name: Optional[str] = "",
     ) -> list[Union[str, int]]:
         """Update data in table. If primary key is duplicated, replace it.
-        
+
         Args:
             collection_name (string): collection name
             data (Union[Dict, List[Dict]]): data that will be upserted

pyobvector/client/ob_client.py

@@ -51,7 +51,9 @@ class ObClient:
         db_name: str = "test",
         **kwargs,
     ):
-        registry.register("mysql.oceanbase", "pyobvector.schema.dialect", "OceanBaseDialect")
+        registry.register(
+            "mysql.oceanbase", "pyobvector.schema.dialect", "OceanBaseDialect"
+        )
 
         setattr(func_mod, "l2_distance", l2_distance)
         setattr(func_mod, "cosine_distance", cosine_distance)
@@ -88,27 +90,31 @@ class ObClient:
             for table_name in tables:
                 if table_name in self.metadata_obj.tables:
                     self.metadata_obj.remove(Table(table_name, self.metadata_obj))
-            self.metadata_obj.reflect(bind=self.engine, only=tables, extend_existing=True)
+            self.metadata_obj.reflect(
+                bind=self.engine, only=tables, extend_existing=True
+            )
         else:
             self.metadata_obj.clear()
             self.metadata_obj.reflect(bind=self.engine, extend_existing=True)
 
     def _is_seekdb(self) -> bool:
         """Check if the database is SeekDB by querying version.
-        
+
         Returns:
             bool: True if database is SeekDB, False otherwise
         """
         is_seekdb = False
         try:
-            if hasattr(self, '_is_seekdb_cached'):
+            if hasattr(self, "_is_seekdb_cached"):
                 return self._is_seekdb_cached
             with self.engine.connect() as conn:
                 result = conn.execute(text("SELECT VERSION()"))
                 version_str = [r[0] for r in result][0]
-                is_seekdb = "SeekDB" in version_str
+                is_seekdb = "seekdb" in version_str.lower()
                 self._is_seekdb_cached = is_seekdb
-                logger.debug(f"Version query result: {version_str}, is_seekdb: {is_seekdb}")
+                logger.debug(
+                    f"Version query result: {version_str}, is_seekdb: {is_seekdb}"
+                )
         except Exception as e:
             logger.warning(f"Failed to query version: {e}")
         return is_seekdb
@@ -414,10 +420,12 @@ class ObClient:
                 if partition_names is None:
                     execute_res = conn.execute(stmt)
                 else:
-                    stmt_str = str(stmt.compile(
-                        dialect=self.engine.dialect,
-                        compile_kwargs={"literal_binds": True}
-                    ))
+                    stmt_str = str(
+                        stmt.compile(
+                            dialect=self.engine.dialect,
+                            compile_kwargs={"literal_binds": True},
+                        )
+                    )
                     stmt_str = self._insert_partition_hint_for_query_sql(
                         stmt_str, f"PARTITION({', '.join(partition_names)})"
                     )
@@ -451,9 +459,7 @@ class ObClient:
 
         with self.engine.connect() as conn:
             with conn.begin():
-                conn.execute(
-                    text(f"ALTER TABLE `{table_name}` {columns_ddl}")
-                )
+                conn.execute(text(f"ALTER TABLE `{table_name}` {columns_ddl}"))
 
         self.refresh_metadata([table_name])
 
@@ -472,8 +478,6 @@ class ObClient:
 
         with self.engine.connect() as conn:
             with conn.begin():
-                conn.execute(
-                    text(f"ALTER TABLE `{table_name}` {columns_ddl}")
-                )
+                conn.execute(text(f"ALTER TABLE `{table_name}` {columns_ddl}"))
 
         self.refresh_metadata([table_name])

pyobvector/client/ob_vec_client.py

@@ -1,4 +1,5 @@
 """OceanBase Vector Store Client."""
+
 import logging
 from typing import Optional, Union
 
@@ -44,18 +45,14 @@ class ObVecClient(ObClient):
         if self.ob_version < ObVersion.from_db_version_nums(4, 3, 3, 0):
             raise ClusterVersionException(
                 code=ErrorCode.NOT_SUPPORTED,
-                message=ExceptionsMessage.ClusterVersionIsLow % ("Vector Store", "4.3.3.0"),
+                message=ExceptionsMessage.ClusterVersionIsLow
+                % ("Vector Store", "4.3.3.0"),
             )
 
-    def _get_sparse_vector_index_params(
-        self, vidxs: Optional[IndexParams]
-    ):
+    def _get_sparse_vector_index_params(self, vidxs: Optional[IndexParams]):
         if vidxs is None:
             return None
-        return [
-            vidx for vidx in vidxs
-            if vidx.is_index_type_sparse_vector()
-        ]
+        return [vidx for vidx in vidxs if vidx.is_index_type_sparse_vector()]
 
     def create_table_with_index_params(
         self,
@@ -100,14 +97,20 @@ class ObVecClient(ObClient):
                     )
                 if sparse_vidxs is not None and len(sparse_vidxs) > 0:
                     create_table_sql = str(CreateTable(table).compile(self.engine))
-                    new_sql = create_table_sql[:create_table_sql.rfind(')')]
+                    # Preserve table options (e.g. ORGANIZATION=heap) after the closing ")"
+                    last_paren = create_table_sql.rfind(")")
+                    table_options_suffix = create_table_sql[
+                        last_paren:
+                    ]  # e.g. ")ORGANIZATION=heap"
+                    new_sql = create_table_sql[:last_paren]
                     for sparse_vidx in sparse_vidxs:
                         sparse_params = sparse_vidx._parse_kwargs()
-                        if 'type' in sparse_params:
+                        if "type" in sparse_params:
                             new_sql += f",\n\tVECTOR INDEX {sparse_vidx.index_name}({sparse_vidx.field_name}) with (type={sparse_params['type']}, distance=inner_product)"
                         else:
                             new_sql += f",\n\tVECTOR INDEX {sparse_vidx.index_name}({sparse_vidx.field_name}) with (distance=inner_product)"
-                    new_sql += "\n)"
+                    # Restore table options after the new closing ")"
+                    new_sql += "\n)" + table_options_suffix[1:]
                     conn.execute(text(new_sql))
                 else:
                     table.create(self.engine, checkfirst=True)
@@ -130,7 +133,9 @@ class ObVecClient(ObClient):
                 # create fts indexes
                 if fts_idxs is not None:
                     for fts_idx in fts_idxs:
-                        idx_cols = [table.c[field_name] for field_name in fts_idx.field_names]
+                        idx_cols = [
+                            table.c[field_name] for field_name in fts_idx.field_names
+                        ]
                         fts_idx = FtsIndex(
                             fts_idx.index_name,
                             fts_idx.param_str(),
@@ -195,7 +200,7 @@ class ObVecClient(ObClient):
         fts_idx_param: FtsIndexParam,
     ):
         """Create fts index with fts index parameter.
-        
+
         Args:
             table_name (string): table name
             fts_idx_param (FtsIndexParam): fts index parameter
@@ -203,7 +208,9 @@ class ObVecClient(ObClient):
         table = Table(table_name, self.metadata_obj, autoload_with=self.engine)
         with self.engine.connect() as conn:
             with conn.begin():
-                idx_cols = [table.c[field_name] for field_name in fts_idx_param.field_names]
+                idx_cols = [
+                    table.c[field_name] for field_name in fts_idx_param.field_names
+                ]
                 fts_idx = FtsIndex(
                     fts_idx_param.index_name,
                     fts_idx_param.param_str(),
@@ -335,11 +342,7 @@ class ObVecClient(ObClient):
                     )
                 )
             else:
-                columns.append(
-                    distance_func(
-                        table.c[vec_column_name], f"{vec_data}"
-                    )
-                )
+                columns.append(distance_func(table.c[vec_column_name], f"{vec_data}"))
         # if idx_name_hint is not None:
         #     stmt = select(*columns).with_hint(
         #         table,
@@ -360,9 +363,7 @@ class ObVecClient(ObClient):
                     "[" + ",".join([str(np.float32(v)) for v in vec_data]) + "]",
                 )
             else:
-                dist_expr = distance_func(
-                    table.c[vec_column_name], f"{vec_data}"
-                )
+                dist_expr = distance_func(table.c[vec_column_name], f"{vec_data}")
             stmt = stmt.where(dist_expr <= distance_threshold)
 
         if isinstance(vec_data, list):
@@ -373,23 +374,23 @@ class ObVecClient(ObClient):
                 )
             )
         else:
-            stmt = stmt.order_by(
-                distance_func(
-                    table.c[vec_column_name], f"{vec_data}"
+            stmt = stmt.order_by(distance_func(table.c[vec_column_name], f"{vec_data}"))
+        stmt_str = (
+            str(
+                stmt.compile(
+                    dialect=self.engine.dialect, compile_kwargs={"literal_binds": True}
                 )
             )
-        stmt_str = (
-            str(stmt.compile(
-                dialect=self.engine.dialect,
-                compile_kwargs={"literal_binds": True}
-            ))
             + f" APPROXIMATE limit {topk}"
         )
         with self.engine.connect() as conn:
             with conn.begin():
                 if idx_name_hint is not None:
                     idx = stmt_str.find("SELECT ")
-                    stmt_str = f"SELECT /*+ index({table_name} {idx_name_hint}) */ " + stmt_str[idx + len("SELECT "):]
+                    stmt_str = (
+                        f"SELECT /*+ index({table_name} {idx_name_hint}) */ "
+                        + stmt_str[idx + len("SELECT ") :]
+                    )
 
                 if partition_names is None:
                     return conn.execute(text(stmt_str))
@@ -433,7 +434,9 @@ class ObVecClient(ObClient):
 
         columns = []
         if output_column_names is not None:
-            columns.extend([table.c[column_name] for column_name in output_column_names])
+            columns.extend(
+                [table.c[column_name] for column_name in output_column_names]
+            )
         else:
             columns.extend([table.c[column.name] for column in table.columns])
         if extra_output_cols is not None:
@@ -462,16 +465,20 @@ class ObVecClient(ObClient):
                 if partition_names is None:
                     if str_list is not None:
                         str_list.append(
-                            str(stmt.compile(
-                                dialect=self.engine.dialect,
-                                compile_kwargs={"literal_binds": True}
-                            ))
+                            str(
+                                stmt.compile(
+                                    dialect=self.engine.dialect,
+                                    compile_kwargs={"literal_binds": True},
+                                )
+                            )
                         )
                     return conn.execute(stmt)
-                stmt_str = str(stmt.compile(
-                    dialect=self.engine.dialect,
-                    compile_kwargs={"literal_binds": True}
-                ))
+                stmt_str = str(
+                    stmt.compile(
+                        dialect=self.engine.dialect,
+                        compile_kwargs={"literal_binds": True},
+                    )
+                )
                 stmt_str = self._insert_partition_hint_for_query_sql(
                     stmt_str, f"PARTITION({', '.join(partition_names)})"
                 )