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

pyobvector/schema/dialect.py

@@ -1,4 +1,5 @@
 """OceanBase dialect."""
+
 from sqlalchemy import util
 from sqlalchemy.dialects.mysql import aiomysql, pymysql
 
@@ -7,10 +8,12 @@ from .vector import VECTOR
 from .sparse_vector import SPARSE_VECTOR
 from .geo_srid_point import POINT
 
+
 class OceanBaseDialect(pymysql.MySQLDialect_pymysql):
     # not change dialect name, since it is a subclass of pymysql.MySQLDialect_pymysql
     # name = "oceanbase"
     """Ocenbase dialect."""
+
     supports_statement_cache = True
 
     def __init__(self, **kwargs):
@@ -36,6 +39,7 @@ class OceanBaseDialect(pymysql.MySQLDialect_pymysql):
 
 class AsyncOceanBaseDialect(aiomysql.MySQLDialect_aiomysql):
     """OceanBase async dialect."""
+
     supports_statement_cache = True
 
     def __init__(self, **kwargs):

pyobvector/schema/full_text_index.py

@@ -1,4 +1,5 @@
 """FullTextIndex: full text search index type"""
+
 from sqlalchemy import Index
 from sqlalchemy.schema import DDLElement
 from sqlalchemy.ext.compiler import compiles
@@ -7,16 +8,18 @@ from sqlalchemy.sql.ddl import SchemaGenerator
 
 class CreateFtsIndex(DDLElement):
     """A new statement clause to create fts index.
-    
+
     Attributes:
     index : fts index schema
     """
+
     def __init__(self, index):
         self.index = index
 
 
 class ObFtsSchemaGenerator(SchemaGenerator):
     """A new schema generator to handle create fts index statement."""
+
     def visit_fts_index(self, index, create_ok=False):
         """Handle create fts index statement compiling.
 
@@ -29,8 +32,10 @@ class ObFtsSchemaGenerator(SchemaGenerator):
         with self.with_ddl_events(index):
             CreateFtsIndex(index)._invoke_with(self.connection)
 
+
 class FtsIndex(Index):
     """Fts Index schema."""
+
     __visit_name__ = "fts_index"
 
     def __init__(self, name, fts_parser: str, *column_names, **kw):
@@ -39,7 +44,7 @@ class FtsIndex(Index):
 
     def create(self, bind, checkfirst: bool = False) -> None:
         """Create fts index.
-        
+
         Args:
             bind: SQL engine or connection.
             checkfirst: check the index exists or not.
@@ -48,7 +53,7 @@ class FtsIndex(Index):
 
 
 @compiles(CreateFtsIndex)
-def compile_create_fts_index(element, compiler, **kw): # pylint: disable=unused-argument
+def compile_create_fts_index(element, compiler, **kw):  # pylint: disable=unused-argument
     """A decorator function to compile create fts index statement."""
     index = element.index
     table_name = index.table.name

pyobvector/schema/geo_srid_point.py

@@ -1,23 +1,26 @@
 """Point: OceanBase GIS data type for SQLAlchemy"""
+
 from typing import Optional
 from sqlalchemy.types import UserDefinedType, String
 
+
 class POINT(UserDefinedType):
     """Point data type definition."""
+
     cache_ok = True
     _string = String()
 
     def __init__(
         self,
         # lat_long: Tuple[float, float],
-        srid: Optional[int] = None
+        srid: Optional[int] = None,
     ):
         """Init Latitude and Longitude."""
         super(UserDefinedType, self).__init__()
         # self.lat_long = lat_long
         self.srid = srid
 
-    def get_col_spec(self, **kw): # pylint: disable=unused-argument
+    def get_col_spec(self, **kw):  # pylint: disable=unused-argument
         """Parse to Point data type definition in text SQL."""
         if self.srid is None:
             return "POINT"

pyobvector/schema/gis_func.py

@@ -10,31 +10,34 @@ from .geo_srid_point import POINT
 
 logger = logging.getLogger(__name__)
 
+
 class ST_GeomFromText(FunctionElement):
     """ST_GeomFromText: parse text to geometry object.
-    
+
     Attributes:
     type : result type
     """
+
     type = BINARY()
 
     def __init__(self, *args):
         super().__init__()
         self.args = args
 
+
 @compiles(ST_GeomFromText)
-def compile_ST_GeomFromText(element, compiler, **kwargs): # pylint: disable=unused-argument
+def compile_ST_GeomFromText(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile ST_GeomFromText function."""
     args = []
     for idx, arg in enumerate(element.args):
         if idx == 0:
             if (
-                (not isinstance(arg, tuple)) or
-                (len(arg) != 2) or
-                (not all(isinstance(x, float) for x in arg))
+                (not isinstance(arg, tuple))
+                or (len(arg) != 2)
+                or (not all(isinstance(x, float) for x in arg))
             ):
                 raise ValueError(
-                    f"Tuple[float, float] is expected for Point literal," \
+                    f"Tuple[float, float] is expected for Point literal,"
                     f"while get {type(arg)}"
                 )
             args.append(f"'{POINT.to_db(arg)}'")
@@ -44,12 +47,14 @@ def compile_ST_GeomFromText(element, compiler, **kwargs): # pylint: disable=unus
     # logger.info(f"{args_str}")
     return f"ST_GeomFromText({args_str})"
 
+
 class st_distance(FunctionElement):
     """st_distance: calculate distance between Points.
-    
+
     Attributes:
     type : result type
     """
+
     type = Float()
     inherit_cache = True
 
@@ -57,19 +62,22 @@ class st_distance(FunctionElement):
         super().__init__()
         self.args = args
 
+
 @compiles(st_distance)
-def compile_st_distance(element, compiler, **kwargs): # pylint: disable=unused-argument
+def compile_st_distance(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile st_distance function."""
     args = ", ".join(compiler.process(arg) for arg in element.args)
     return f"st_distance({args})"
 
+
 class st_dwithin(FunctionElement):
     """st_dwithin: Checks if the distance between two points
     is less than a specified distance.
-    
+
     Attributes:
     type : result type
     """
+
     type = Boolean()
     inherit_cache = True
 
@@ -77,8 +85,9 @@ class st_dwithin(FunctionElement):
         super().__init__()
         self.args = args
 
+
 @compiles(st_dwithin)
-def compile_st_dwithin(element, compiler, **kwargs): # pylint: disable=unused-argument
+def compile_st_dwithin(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile st_dwithin function."""
     args = []
     for idx, arg in enumerate(element.args):
@@ -89,12 +98,14 @@ def compile_st_dwithin(element, compiler, **kwargs): # pylint: disable=unused-ar
     args_str = ", ".join(args)
     return f"_st_dwithin({args_str})"
 
+
 class st_astext(FunctionElement):
     """st_astext: Returns a Point in human-readable format.
 
     Attributes:
     type : result type
     """
+
     type = Text()
     inherit_cache = True
 
@@ -102,8 +113,9 @@ class st_astext(FunctionElement):
         super().__init__()
         self.args = args
 
+
 @compiles(st_astext)
-def compile_st_astext(element, compiler, **kwargs): # pylint: disable=unused-argument
+def compile_st_astext(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile st_astext function."""
     args = ", ".join(compiler.process(arg) for arg in element.args)
     return f"st_astext({args})"

pyobvector/schema/match_against_func.py

@@ -8,31 +8,36 @@ from sqlalchemy import literal, column
 
 logger = logging.getLogger(__name__)
 
+
 class MatchAgainst(FunctionElement):
     """MatchAgainst: match clause for full text search.
 
     Attributes:
     type : result type
     """
+
     inherit_cache = True
 
     def __init__(self, query, *columns):
         columns = [column(col) if isinstance(col, str) else col for col in columns]
         super().__init__(literal(query), *columns)
-    
+
+
 @compiles(MatchAgainst)
-def complie_MatchAgainst(element, compiler, **kwargs): # pylint: disable=unused-argument
+def complie_MatchAgainst(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile MatchAgainst function."""
     clauses = list(element.clauses)
     if len(clauses) < 2:
         raise ValueError(
-            f"MatchAgainst should take a string expression and " \
+            f"MatchAgainst should take a string expression and "
             f"at least one column name string as parameters."
         )
-    
+
     query_expr = clauses[0]
     compiled_query = compiler.process(query_expr, **kwargs)
     column_exprs = clauses[1:]
-    compiled_columns = [compiler.process(col, identifier_prepared=True) for col in column_exprs]
+    compiled_columns = [
+        compiler.process(col, identifier_prepared=True) for col in column_exprs
+    ]
     columns_str = ", ".join(compiled_columns)
     return f"MATCH ({columns_str}) AGAINST ({compiled_query} IN NATURAL LANGUAGE MODE)"

pyobvector/schema/ob_table.py

@@ -1,9 +1,11 @@
 """ObTable: extension to Table for creating table with vector index."""
+
 from sqlalchemy import Table
 from .vector_index import ObSchemaGenerator
 
 
 class ObTable(Table):
     """A class extends SQLAlchemy Table to do table creation with vector index."""
+
     def create(self, bind, checkfirst: bool = False) -> None:
         bind._run_ddl_visitor(ObSchemaGenerator, self, checkfirst=checkfirst)

pyobvector/schema/reflection.py

@@ -1,14 +1,21 @@
 """OceanBase table definition reflection."""
+
 import re
 import logging
-from sqlalchemy.dialects.mysql.reflection import MySQLTableDefinitionParser, _re_compile, cleanup_text
+from sqlalchemy.dialects.mysql.reflection import (
+    MySQLTableDefinitionParser,
+    _re_compile,
+    cleanup_text,
+)
 
 from pyobvector.schema.array import ARRAY
 
 logger = logging.getLogger(__name__)
 
+
 class OceanBaseTableDefinitionParser(MySQLTableDefinitionParser):
     """OceanBase table definition parser."""
+
     def __init__(self, dialect, preparer, *, default_schema=None):
         MySQLTableDefinitionParser.__init__(self, dialect, preparer)
         self.default_schema = default_schema
@@ -82,15 +89,19 @@ class OceanBaseTableDefinitionParser(MySQLTableDefinitionParser):
         m = self._re_array_column.match(line)
         if m:
             spec = m.groupdict()
-            name, coltype_with_args = spec["name"].strip(), spec["coltype_with_args"].strip()
+            name, coltype_with_args = (
+                spec["name"].strip(),
+                spec["coltype_with_args"].strip(),
+            )
 
             item_pattern = re.compile(
-                r"^(?:array\s*\()*([\w]+)(?:\(([\d,]+)\))?\)*$",
-                re.IGNORECASE
+                r"^(?:array\s*\()*([\w]+)(?:\(([\d,]+)\))?\)*$", re.IGNORECASE
             )
             item_m = item_pattern.match(coltype_with_args)
             if not item_m:
-                raise ValueError(f"Failed to find inner type from array column definition: {line}")
+                raise ValueError(
+                    f"Failed to find inner type from array column definition: {line}"
+                )
 
             item_type = self.dialect.ischema_names[item_m.group(1).lower()]
             item_type_arg = item_m.group(2)
@@ -99,9 +110,11 @@ class OceanBaseTableDefinitionParser(MySQLTableDefinitionParser):
             elif item_type_arg[0] == "'" and item_type_arg[-1] == "'":
                 item_type_args = self._re_csv_str.findall(item_type_arg)
             else:
-                item_type_args = [int(v) for v in self._re_csv_int.findall(item_type_arg)]
+                item_type_args = [
+                    int(v) for v in self._re_csv_int.findall(item_type_arg)
+                ]
 
-            nested_level = coltype_with_args.lower().count('array')
+            nested_level = coltype_with_args.lower().count("array")
             type_instance = item_type(*item_type_args)
             for _ in range(nested_level):
                 type_instance = ARRAY(type_instance)
@@ -144,7 +157,11 @@ class OceanBaseTableDefinitionParser(MySQLTableDefinitionParser):
 
             if tp == "fk_constraint":
                 table = spec.get("table", [])
-                if isinstance(table, list) and len(table) == 2 and table[0] == self.default_schema:
+                if (
+                    isinstance(table, list)
+                    and len(table) == 2
+                    and table[0] == self.default_schema
+                ):
                     spec["table"] = table[1:]
 
             for action in ["onupdate", "ondelete"]:

pyobvector/schema/replace_stmt.py

@@ -1,11 +1,15 @@
 """ReplaceStmt: replace into statement compilation."""
+
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.sql.expression import Insert
 
+
 class ReplaceStmt(Insert):
     """Replace into statement."""
+
     inherit_cache = True
 
+
 @compiles(ReplaceStmt)
 def compile_replace_stmt(insert, compiler, **kw):
     """Compile replace into statement.

pyobvector/schema/sparse_vector.py

@@ -1,25 +1,28 @@
 """SPARSE_VECTOR: An extended data type for SQLAlchemy"""
+
 from sqlalchemy.types import UserDefinedType, String
 from ..util import SparseVector
 
+
 class SPARSE_VECTOR(UserDefinedType):
     """SPARSE_VECTOR data type definition."""
+
     cache_ok = True
     _string = String()
 
     def __init__(self):
         super(UserDefinedType, self).__init__()
 
-    def get_col_spec(self, **kw): # pylint: disable=unused-argument
+    def get_col_spec(self, **kw):  # pylint: disable=unused-argument
         """Parse to sparse vector data type definition in text SQL."""
         return "SPARSEVECTOR"
-    
+
     def bind_processor(self, dialect):
         def process(value):
             return SparseVector._to_db(value)
 
         return process
-    
+
     def literal_processor(self, dialect):
         string_literal_processor = self._string._cached_literal_processor(dialect)
 
@@ -32,4 +35,4 @@ class SPARSE_VECTOR(UserDefinedType):
         def process(value):
             return SparseVector._from_db(value)
 
-        return process
+        return process

pyobvector/schema/vec_dist_func.py

@@ -11,6 +11,7 @@ from sqlalchemy import Float
 
 logger = logging.getLogger(__name__)
 
+
 def parse_vec_distance_func_args(element, compiler, **kwargs):
     args = []
     for arg in element.args:
@@ -21,20 +22,23 @@ def parse_vec_distance_func_args(element, compiler, **kwargs):
     args = ", ".join(args)
     return args
 
+
 class l2_distance(FunctionElement):
     """Vector distance function: l2_distance.
-    
+
     Attributes:
     type : result type
     """
+
     type = Float()
 
     def __init__(self, *args):
         super().__init__()
         self.args = args
 
+
 @compiles(l2_distance)
-def compile_l2_distance(element, compiler, **kwargs): # pylint: disable=unused-argument
+def compile_l2_distance(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile l2_distance function.
 
     Args:
@@ -46,41 +50,47 @@ def compile_l2_distance(element, compiler, **kwargs): # pylint: disable=unused-a
 
 class cosine_distance(FunctionElement):
     """Vector distance function: cosine_distance.
-    
+
     Attributes:
     type : result type
     """
+
     type = Float()
 
     def __init__(self, *args):
         super().__init__()
         self.args = args
 
+
 @compiles(cosine_distance)
-def compile_cosine_distance(element, compiler, **kwargs): # pylint: disable=unused-argument
+def compile_cosine_distance(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile cosine_distance function.
 
     Args:
         element: cosine_distance arguments
         compiler: SQL compiler
     """
-    return f"cosine_distance({parse_vec_distance_func_args(element, compiler, **kwargs)})"
+    return (
+        f"cosine_distance({parse_vec_distance_func_args(element, compiler, **kwargs)})"
+    )
 
 
 class inner_product(FunctionElement):
     """Vector distance function: inner_product.
-    
+
     Attributes:
     type : result type
     """
+
     type = Float()
 
     def __init__(self, *args):
         super().__init__()
         self.args = args
 
+
 @compiles(inner_product)
-def compile_inner_product(element, compiler, **kwargs): # pylint: disable=unused-argument
+def compile_inner_product(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile inner_product function.
 
     Args:
@@ -89,20 +99,23 @@ def compile_inner_product(element, compiler, **kwargs): # pylint: disable=unused
     """
     return f"inner_product({parse_vec_distance_func_args(element, compiler, **kwargs)})"
 
+
 class negative_inner_product(FunctionElement):
     """Vector distance function: negative_inner_product.
-    
+
     Attributes:
     type : result type
     """
+
     type = Float()
 
     def __init__(self, *args):
         super().__init__()
         self.args = args
 
+
 @compiles(negative_inner_product)
-def compile_negative_inner_product(element, compiler, **kwargs): # pylint: disable=unused-argument
+def compile_negative_inner_product(element, compiler, **kwargs):  # pylint: disable=unused-argument
     """Compile negative_inner_product function.
 
     Args:

pyobvector/schema/vector.py

@@ -1,10 +1,12 @@
 """VECTOR: An extended data type for SQLAlchemy"""
+
 from sqlalchemy.types import UserDefinedType, String
 from ..util import Vector
 
 
 class VECTOR(UserDefinedType):
     """VECTOR data type definition."""
+
     cache_ok = True
     _string = String()
 
@@ -12,7 +14,7 @@ class VECTOR(UserDefinedType):
         super(UserDefinedType, self).__init__()
         self.dim = dim
 
-    def get_col_spec(self, **kw): # pylint: disable=unused-argument
+    def get_col_spec(self, **kw):  # pylint: disable=unused-argument
         """Parse to vector data type definition in text SQL."""
         if self.dim is None:
             return "VECTOR"

pyobvector/schema/vector_index.py

@@ -1,4 +1,5 @@
 """VectorIndex: An extended index type for SQLAlchemy"""
+
 from sqlalchemy import Index
 from sqlalchemy.schema import DDLElement
 from sqlalchemy.ext.compiler import compiles
@@ -7,16 +8,18 @@ from sqlalchemy.sql.ddl import SchemaGenerator
 
 class CreateVectorIndex(DDLElement):
     """A new statement clause to create vector index.
-    
+
     Attributes:
         index: vector index schema
     """
+
     def __init__(self, index):
         self.index = index
 
 
 class ObSchemaGenerator(SchemaGenerator):
     """A new schema generator to handle create vector index statement."""
+
     def visit_vector_index(self, index, create_ok=False):
         """Handle create vector index statement compiling.
 
@@ -32,6 +35,7 @@ class ObSchemaGenerator(SchemaGenerator):
 
 class VectorIndex(Index):
     """Vector Index schema."""
+
     __visit_name__ = "vector_index"
 
     def __init__(self, name, *column_names, params: str = None, **kw):
@@ -44,7 +48,7 @@ class VectorIndex(Index):
 
     def create(self, bind, checkfirst: bool = False) -> None:
         """Create vector index.
-        
+
         Args:
             bind: SQL engine or connection.
             checkfirst: check the index exists or not.
@@ -53,7 +57,7 @@ class VectorIndex(Index):
 
 
 @compiles(CreateVectorIndex)
-def compile_create_vector_index(element, compiler, **kw): # pylint: disable=unused-argument
+def compile_create_vector_index(element, compiler, **kw):  # pylint: disable=unused-argument
     """A decorator function to compile create vector index statement."""
     index = element.index
     table_name = index.table.name

pyobvector/util/__init__.py

@@ -4,6 +4,7 @@
 * SparseVector  A utility class for the extended data type class 'SPARSE_VECTOR'
 * ObVersion OceanBase cluster version class
 """
+
 from .vector import Vector
 from .sparse_vector import SparseVector
 from .ob_version import ObVersion

pyobvector/util/ob_version.py

@@ -1,4 +1,5 @@
 """OceanBase cluster version module."""
+
 import copy
 
 
@@ -8,6 +9,7 @@ class ObVersion:
     Attributes:
         version_nums (List[int]): version number of OceanBase cluster. For example, '4.3.3.0'
     """
+
     def __init__(self, version_nums: list[int]):
         self.version_nums = copy.deepcopy(version_nums)
 

pyobvector/util/sparse_vector.py

@@ -1,24 +1,27 @@
 """A utility module for the extended data type class 'SPARSE_VECTOR'."""
+
 import ast
 
+
 class SparseVector:
     """A transformer class between python dict and OceanBase SPARSE_VECTOR.
 
     Attributes:
     _value (Dict) : a python dict
     """
+
     def __init__(self, value):
         if not isinstance(value, dict):
             raise ValueError("Sparse Vector should be a dict in python")
-        
+
         self._value = value
-    
+
     def __repr__(self):
         return f"{self._value}"
-    
+
     def to_text(self):
         return f"{self._value}"
-    
+
     @classmethod
     def from_text(cls, value: str):
         """Construct Sparse Vector class with dict in string format.
@@ -27,7 +30,7 @@ class SparseVector:
             value: For example, '{1:1.1, 2:2.2}'
         """
         return cls(ast.literal_eval(value))
-    
+
     @classmethod
     def _to_db(cls, value):
         if value is None:
@@ -45,4 +48,4 @@ class SparseVector:
 
         if isinstance(value, str):
             return cls.from_text(value)._value
-        raise ValueError(f"unexpected sparse vector type: {type(value)}")
+        raise ValueError(f"unexpected sparse vector type: {type(value)}")

pyobvector/util/vector.py

@@ -1,4 +1,5 @@
 """A utility module for the extended data type class 'VECTOR'."""
+
 import json
 import numpy as np
 
@@ -9,6 +10,7 @@ class Vector:
     Attributes:
         _value (numpy.array): a numpy array
     """
+
     def __init__(self, value):
         # big-endian float32
         if not isinstance(value, np.ndarray) or value.dtype != ">f4":