sqlalchemy-cratedb 0.36.1__py2.py3-none-any.whl → 0.37.0__py2.py3-none-any.whl

sqlalchemy_cratedb/__init__.py

@@ -27,6 +27,7 @@ from .support import insert_bulk
 from .type.array import ObjectArray
 from .type.geo import Geopoint, Geoshape
 from .type.object import ObjectType
+from .type.vector import FloatVector, knn_match
 
 if SA_VERSION < SA_1_4:
     import textwrap
@@ -51,9 +52,11 @@ if SA_VERSION < SA_1_4:
 
 __all__ = [
     dialect,
+    FloatVector,
     Geopoint,
     Geoshape,
     ObjectArray,
     ObjectType,
     match,
+    knn_match,
 ]

sqlalchemy_cratedb/compiler.py

@@ -238,6 +238,12 @@ class CrateTypeCompiler(compiler.GenericTypeCompiler):
     def visit_OBJECT(self, type_, **kw):
         return "OBJECT"
 
+    def visit_FLOAT_VECTOR(self, type_, **kw):
+        dimensions = type_.dimensions
+        if dimensions is None:
+            raise ValueError("FloatVector must be initialized with dimension size")
+        return f"FLOAT_VECTOR({dimensions})"
+
 
 class CrateCompiler(compiler.SQLCompiler):
 

sqlalchemy_cratedb/dialect.py

@@ -33,7 +33,7 @@ from .compiler import (
 )
 from crate.client.exceptions import TimezoneUnawareException
 from .sa_version import SA_VERSION, SA_1_4, SA_2_0
-from .type import ObjectArray, ObjectType
+from .type import FloatVector, ObjectArray, ObjectType
 
 TYPES_MAP = {
     "boolean": sqltypes.Boolean,
@@ -51,7 +51,8 @@ TYPES_MAP = {
     "float": sqltypes.Float,
     "real": sqltypes.Float,
     "string": sqltypes.String,
-    "text": sqltypes.String
+    "text": sqltypes.String,
+    "float_vector": FloatVector,
 }
 try:
     # SQLAlchemy >= 1.1
@@ -228,6 +229,15 @@ class CrateDialect(default.DefaultDialect):
     def _get_default_schema_name(self, connection):
         return 'doc'
 
+    def _get_effective_schema_name(self, connection):
+        schema_name_raw = connection.engine.url.query.get("schema")
+        schema_name = None
+        if isinstance(schema_name_raw, str):
+            schema_name = schema_name_raw
+        elif isinstance(schema_name_raw, tuple):
+            schema_name = schema_name_raw[0]
+        return schema_name
+
     def _get_server_version_info(self, connection):
         return tuple(connection.connection.lowest_server_version.version)
 
@@ -257,6 +267,8 @@ class CrateDialect(default.DefaultDialect):
 
     @reflection.cache
     def get_table_names(self, connection, schema=None, **kw):
+        if schema is None:
+            schema = self._get_effective_schema_name(connection)
         cursor = connection.exec_driver_sql(
             "SELECT table_name FROM information_schema.tables "
             "WHERE {0} = ? "

sqlalchemy_cratedb/type/__init__.py

@@ -1,3 +1,4 @@
 from .array import ObjectArray
 from .geo import Geopoint, Geoshape
 from .object import ObjectType
+from .vector import FloatVector, knn_match

sqlalchemy_cratedb/type/vector.py

@@ -0,0 +1,173 @@
+"""
+## About
+SQLAlchemy data type implementation for CrateDB's `FLOAT_VECTOR` type.
+
+## References
+- https://cratedb.com/docs/crate/reference/en/latest/general/ddl/data-types.html#float-vector
+- https://cratedb.com/docs/crate/reference/en/latest/general/builtins/scalar-functions.html#scalar-knn-match
+
+## Details
+The implementation is based on SQLAlchemy's `TypeDecorator`, and also
+offers compiler support.
+
+## Notes
+CrateDB currently only supports the similarity function `VectorSimilarityFunction.EUCLIDEAN`.
+-- https://github.com/crate/crate/blob/5.5.1/server/src/main/java/io/crate/types/FloatVectorType.java#L55
+
+pgvector use a comparator to apply different similarity functions as operators,
+see `pgvector.sqlalchemy.Vector.comparator_factory`.
+
+<->: l2/euclidean_distance
+<#>: max_inner_product
+<=>: cosine_distance
+
+## Backlog
+- After dropping support for SQLAlchemy 1.3, use
+  `class FloatVector(sa.TypeDecorator[t.Sequence[float]]):`
+
+## Origin
+This module is based on the corresponding pgvector implementation
+by Andrew Kane. Thank you.
+
+The MIT License (MIT)
+Copyright (c) 2021-2023 Andrew Kane
+https://github.com/pgvector/pgvector-python
+"""
+import typing as t
+
+if t.TYPE_CHECKING:
+    import numpy.typing as npt  # pragma: no cover
+
+import sqlalchemy as sa
+from sqlalchemy.sql.expression import ColumnElement, literal
+from sqlalchemy.ext.compiler import compiles
+
+
+__all__ = [
+    "from_db",
+    "knn_match",
+    "to_db",
+    "FloatVector",
+]
+
+
+def from_db(value: t.Iterable) -> t.Optional["npt.ArrayLike"]:
+    import numpy as np
+
+    # from `pgvector.utils`
+    # could be ndarray if already cast by lower-level driver
+    if value is None or isinstance(value, np.ndarray):
+        return value
+
+    return np.array(value, dtype=np.float32)
+
+
+def to_db(value: t.Any, dim: t.Optional[int] = None) -> t.Optional[t.List]:
+    import numpy as np
+
+    # from `pgvector.utils`
+    if value is None:
+        return value
+
+    if isinstance(value, np.ndarray):
+        if value.ndim != 1:
+            raise ValueError("expected ndim to be 1")
+
+        if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating):
+            raise ValueError("dtype must be numeric")
+
+        value = value.tolist()
+
+    if dim is not None and len(value) != dim:
+        raise ValueError("expected %d dimensions, not %d" % (dim, len(value)))
+
+    return value
+
+
+class FloatVector(sa.TypeDecorator):
+    """
+    SQLAlchemy `FloatVector` data type for CrateDB.
+    """
+
+    cache_ok = False
+
+    __visit_name__ = "FLOAT_VECTOR"
+
+    _is_array = True
+
+    zero_indexes = False
+
+    impl = sa.ARRAY
+
+    def __init__(self, dimensions: int = None):
+        super().__init__(sa.FLOAT, dimensions=dimensions)
+
+    def as_generic(self, allow_nulltype=False):
+        return sa.ARRAY(item_type=sa.FLOAT)
+
+    @property
+    def python_type(self):
+        return list
+
+    def bind_processor(self, dialect: sa.engine.Dialect) -> t.Callable:
+        def process(value: t.Iterable) -> t.Optional[t.List]:
+            return to_db(value, self.dimensions)
+
+        return process
+
+    def result_processor(self, dialect: sa.engine.Dialect, coltype: t.Any) -> t.Callable:
+        def process(value: t.Any) -> t.Optional["npt.ArrayLike"]:
+            return from_db(value)
+
+        return process
+
+
+class KnnMatch(ColumnElement):
+    """
+    Wrap CrateDB's `KNN_MATCH` function into an SQLAlchemy function.
+
+    https://cratedb.com/docs/crate/reference/en/latest/general/builtins/scalar-functions.html#scalar-knn-match
+    """
+    inherit_cache = True
+
+    def __init__(self, column, term, k=None):
+        super().__init__()
+        self.column = column
+        self.term = term
+        self.k = k
+
+    def compile_column(self, compiler):
+        return compiler.process(self.column)
+
+    def compile_term(self, compiler):
+        return compiler.process(literal(self.term))
+
+    def compile_k(self, compiler):
+        return compiler.process(literal(self.k))
+
+
+def knn_match(column, term, k):
+    """
+    Generate a match predicate for vector search.
+
+    :param column: A reference to a column or an index.
+
+    :param term: The term to match against. This is an array of floating point
+     values, which is compared to other vectors using a HNSW index search.
+
+    :param k: The `k` argument determines the number of nearest neighbours to
+    search in the index.
+    """
+    return KnnMatch(column, term, k)
+
+
+@compiles(KnnMatch)
+def compile_knn_match(knn_match, compiler, **kwargs):
+    """
+    Clause compiler for `KNN_MATCH`.
+    """
+    return "KNN_MATCH(%s, %s, %s)" % (
+        knn_match.compile_column(compiler),
+        knn_match.compile_term(compiler),
+        knn_match.compile_k(compiler),
+    )

{sqlalchemy_cratedb-0.36.1.dist-info → sqlalchemy_cratedb-0.37.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlalchemy-cratedb
-Version: 0.36.1
+Version: 0.37.0
 Summary: SQLAlchemy dialect for CrateDB.
 Author-email: "Crate.io" <office@crate.io>
 License: Apache License 2.0
@@ -68,6 +68,8 @@ Requires-Dist: sqlalchemy <2.1,>=1
 Requires-Dist: verlib2 ==0.2
 Requires-Dist: backports.zoneinfo <1 ; python_version < "3.9"
 Requires-Dist: importlib-resources ; python_version < "3.9"
+Provides-Extra: all
+Requires-Dist: sqlalchemy-cratedb[vector] ; extra == 'all'
 Provides-Extra: develop
 Requires-Dist: black <25 ; extra == 'develop'
 Requires-Dist: mypy <1.11 ; extra == 'develop'
@@ -82,12 +84,15 @@ Provides-Extra: release
 Requires-Dist: build <2 ; extra == 'release'
 Requires-Dist: twine <6 ; extra == 'release'
 Provides-Extra: test
+Requires-Dist: cratedb-toolkit[testing] ; extra == 'test'
 Requires-Dist: dask[dataframe] ; extra == 'test'
 Requires-Dist: pandas <2.3 ; extra == 'test'
 Requires-Dist: pueblo >=0.0.7 ; extra == 'test'
 Requires-Dist: pytest <9 ; extra == 'test'
 Requires-Dist: pytest-cov <6 ; extra == 'test'
 Requires-Dist: pytest-mock <4 ; extra == 'test'
+Provides-Extra: vector
+Requires-Dist: numpy ; extra == 'vector'
 
 # SQLAlchemy dialect for CrateDB
 

{sqlalchemy_cratedb-0.36.1.dist-info → sqlalchemy_cratedb-0.37.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
-sqlalchemy_cratedb/__init__.py,sha256=ahkoO7im4phxQSceEfW2cybgCCWPpBj58sT6eg3aNbU,2209
-sqlalchemy_cratedb/compiler.py,sha256=FHIq7BXpOZTB-vHm7DkVXmAQ2_D4zYX1nQg-5q0KZ1s,11393
-sqlalchemy_cratedb/dialect.py,sha256=lCZerHVNnh7VuLzq0pFPLRfmsqgplOwXCORoXHrH22Q,13708
+sqlalchemy_cratedb/__init__.py,sha256=HIsl7K75Yi2Z1rey-LXPiuonRVJ-Qj8TVWKmIk5d8g8,2289
+sqlalchemy_cratedb/compiler.py,sha256=SDLeuwmk31KWYrNCXKtBqqatIA3z42GAopRhQLHO3o8,11639
+sqlalchemy_cratedb/dialect.py,sha256=kC0HqTzhPSj9eD2IlqDxJR_1iLL58nEpzKqGY2yETPA,14205
 sqlalchemy_cratedb/predicate.py,sha256=HT7tHF7PxX71G8m91lqpH732iVQZtO8_NPZUIsllcys,3622
 sqlalchemy_cratedb/sa_version.py,sha256=NhNA5FUw0Sm5n-fSVo6CTcaioDRovUknFOslX31ir-I,1168
 sqlalchemy_cratedb/support.py,sha256=gw3I4UUe2pfpfxeyJka-eETTP_2Dz9LkjmgONnO382E,2589
@@ -9,14 +9,15 @@ sqlalchemy_cratedb/compat/api13.py,sha256=koqqq48UmAmrjXOyc5lV7gaVEl8xmi5xHUqTBV
 sqlalchemy_cratedb/compat/core10.py,sha256=P46_ABaii-LQrzeeu8bkhCkoYLr2sVZB3TMGr2BSqD4,8636
 sqlalchemy_cratedb/compat/core14.py,sha256=xh_wnNkMUD8aE_VFQPnPeXWsPIFzPg0ZjkhfJBzo3IM,11606
 sqlalchemy_cratedb/compat/core20.py,sha256=QgNEA4z37ldFI4X9HapbTcga7KvsrdMkoMRjxvCJ4tE,15145
-sqlalchemy_cratedb/type/__init__.py,sha256=iTGuaDZ40LiRXgEVd9RfUS5sODB0wqkMQVZMhuO8zyI,98
+sqlalchemy_cratedb/type/__init__.py,sha256=8Bp0RKTFAx-3ETOcG4R4wjhL5rarN7tJLfrue7S6mH8,141
 sqlalchemy_cratedb/type/array.py,sha256=RrHkFoNmTKtk_6md0OtsR7xofRxdk0PbfO-JFfevBJI,4396
 sqlalchemy_cratedb/type/geo.py,sha256=9wFGxGMxxN-7_qZTlmWialzaPp_zzFj4U-yVqY_deA0,1377
 sqlalchemy_cratedb/type/object.py,sha256=-bebiW38vor3grD9qsomzJ_z3zRpNMb2XLLto2fFez4,3016
-sqlalchemy_cratedb-0.36.1.dist-info/LICENSE,sha256=s_w3FXmAYQuatqsgvyYLnGyC_13KOqp3W1DUEXO9RpY,10175
-sqlalchemy_cratedb-0.36.1.dist-info/METADATA,sha256=oAJheA4F2rvQ0FBLWW95VGMZttyYpJLjmDrd96iYLew,6360
-sqlalchemy_cratedb-0.36.1.dist-info/NOTICE,sha256=yU9CWOf_XrVU7fpqGgM9tDjppoMyfHHBmFVMiINZk-k,1167
-sqlalchemy_cratedb-0.36.1.dist-info/WHEEL,sha256=DZajD4pwLWue70CAfc7YaxT1wLUciNBvN_TTcvXpltE,110
-sqlalchemy_cratedb-0.36.1.dist-info/entry_points.txt,sha256=c14wyCG3OeM64_DUbI_vLVUXR3e3GhDyO_PCjo6UQMU,57
-sqlalchemy_cratedb-0.36.1.dist-info/top_level.txt,sha256=UjjXz0burl_-2MApzLzffHG_2RXm6KljZvoGJHISMPo,19
-sqlalchemy_cratedb-0.36.1.dist-info/RECORD,,
+sqlalchemy_cratedb/type/vector.py,sha256=5Q2v-RuiNKriSm-EX7tTb1PelqqoBvjDOFDB27Xxk5I,4723
+sqlalchemy_cratedb-0.37.0.dist-info/LICENSE,sha256=s_w3FXmAYQuatqsgvyYLnGyC_13KOqp3W1DUEXO9RpY,10175
+sqlalchemy_cratedb-0.37.0.dist-info/METADATA,sha256=JILkri_5NoeIpVpTpZXxczffrmSH50el_VCp0jvEQWk,6561
+sqlalchemy_cratedb-0.37.0.dist-info/NOTICE,sha256=yU9CWOf_XrVU7fpqGgM9tDjppoMyfHHBmFVMiINZk-k,1167
+sqlalchemy_cratedb-0.37.0.dist-info/WHEEL,sha256=DZajD4pwLWue70CAfc7YaxT1wLUciNBvN_TTcvXpltE,110
+sqlalchemy_cratedb-0.37.0.dist-info/entry_points.txt,sha256=c14wyCG3OeM64_DUbI_vLVUXR3e3GhDyO_PCjo6UQMU,57
+sqlalchemy_cratedb-0.37.0.dist-info/top_level.txt,sha256=UjjXz0burl_-2MApzLzffHG_2RXm6KljZvoGJHISMPo,19
+sqlalchemy_cratedb-0.37.0.dist-info/RECORD,,