sqlalchemy-cratedb 0.36.0__tar.gz → 0.37.0__tar.gz

{sqlalchemy_cratedb-0.36.0 → sqlalchemy_cratedb-0.37.0}/CHANGES.md

@@ -3,6 +3,19 @@
 
 ## Unreleased
 
+## 2024/06/13 0.37.0
+- Added support for CrateDB's [FLOAT_VECTOR] data type and its accompanying
+  [KNN_MATCH] function, for HNSW matches. For SQLAlchemy column definitions,
+  you can use it like `FloatVector(dimensions=1536)`.
+- Fixed `get_table_names()` reflection method to respect the
+  `schema` query argument in SQLAlchemy connection URLs.
+
+[FLOAT_VECTOR]: https://cratedb.com/docs/crate/reference/en/latest/general/ddl/data-types.html#float-vector
+[KNN_MATCH]: https://cratedb.com/docs/crate/reference/en/latest/general/builtins/scalar-functions.html#scalar-knn-match
+
+## 2024/06/11 0.36.1
+- Dependencies: Use `crate==1.0.0dev0`
+
 ## 2024/06/11 0.36.0
 - Dependencies: Use `dask[dataframe]`
 - Maintenance release after splitting packages `crate-python` vs.

{sqlalchemy_cratedb-0.36.0 → sqlalchemy_cratedb-0.37.0}/DEVELOP.md

@@ -16,6 +16,7 @@ further commands.
 
 Verify code by running all linters and software tests:
 
+    export CRATEDB_VERSION=latest
     docker compose -f tests/docker-compose.yml up
     poe check
 
@@ -36,25 +37,23 @@ Format code:
 
 On branch `main`:
 
--   Add a section for the new version in the `CHANGES.md` file
--   Commit your changes with a message like \"Release x.y.z\"
--   Push to origin/\<release_branch\>
+-   Add a section for the new version in the `CHANGES.md` file.
+-   Commit your changes with a message like `Release x.y.z`.
 -   Create a tag, and push to remote.
+    This will trigger a GitHub action which releases the new version to PyPi.
     ```shell
-    git tag v0.0.1
+    git tag 0.36.1
     git push --tags
     ```
--   Build package, and upload to PyPI.
-    ```shell
-    poe release
-    ```
-
-On GitHub:
-
--   Designate a new release on GitHub, copying in the relevant section
+-   On GitHub, designate a new release, copying in the relevant section
     from the CHANGELOG.
     https://github.com/crate-workbench/sqlalchemy-cratedb/releases
 
+Optionally, build the package and upload to PyPI manually.
+```shell
+poe release
+```
+
 
 ## Writing documentation
 

{sqlalchemy_cratedb-0.36.0/src/sqlalchemy_cratedb.egg-info → sqlalchemy_cratedb-0.37.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlalchemy-cratedb
-Version: 0.36.0
+Version: 0.37.0
 Summary: SQLAlchemy dialect for CrateDB.
 Author-email: "Crate.io" <office@crate.io>
 License: Apache License 2.0
@@ -63,11 +63,13 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: backports.zoneinfo<1; python_version < "3.9"
-Requires-Dist: crate
+Requires-Dist: crate==1.0.0dev0
 Requires-Dist: geojson<4,>=2.5
 Requires-Dist: importlib-resources; python_version < "3.9"
 Requires-Dist: sqlalchemy<2.1,>=1
 Requires-Dist: verlib2==0.2
+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.0 → sqlalchemy_cratedb-0.37.0}/pyproject.toml

@@ -85,13 +85,16 @@ dynamic = [
 ]
 dependencies = [
   'backports.zoneinfo<1; python_version < "3.9"',
-  "crate",
+  "crate==1.0.0dev0",
   "geojson<4,>=2.5",
   'importlib-resources; python_version < "3.9"',
   "sqlalchemy<2.1,>=1",
   "verlib2==0.2",
 ]
 [project.optional-dependencies]
+all = [
+  "sqlalchemy-cratedb[vector]",
+]
 develop = [
   "black<25",
   "mypy<1.11",
@@ -109,6 +112,7 @@ release = [
   "twine<6",
 ]
 test = [
+  "cratedb-toolkit[testing]",
   "dask[dataframe]",
   "pandas<2.3",
   "pueblo>=0.0.7",
@@ -116,6 +120,9 @@ test = [
   "pytest-cov<6",
   "pytest-mock<4",
 ]
+vector = [
+  "numpy",
+]
 [project.urls]
 changelog = "https://github.com/crate-workbench/sqlalchemy-cratedb/blob/main/CHANGES.md"
 documentation = "https://github.com/crate-workbench/sqlalchemy-cratedb"

{sqlalchemy_cratedb-0.36.0 → sqlalchemy_cratedb-0.37.0}/src/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-0.36.0 → sqlalchemy_cratedb-0.37.0}/src/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-0.36.0 → sqlalchemy_cratedb-0.37.0}/src/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-0.36.0 → sqlalchemy_cratedb-0.37.0}/src/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-0.37.0/src/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.0 → sqlalchemy_cratedb-0.37.0/src/sqlalchemy_cratedb.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlalchemy-cratedb
-Version: 0.36.0
+Version: 0.37.0
 Summary: SQLAlchemy dialect for CrateDB.
 Author-email: "Crate.io" <office@crate.io>
 License: Apache License 2.0
@@ -63,11 +63,13 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: NOTICE
 Requires-Dist: backports.zoneinfo<1; python_version < "3.9"
-Requires-Dist: crate
+Requires-Dist: crate==1.0.0dev0
 Requires-Dist: geojson<4,>=2.5
 Requires-Dist: importlib-resources; python_version < "3.9"
 Requires-Dist: sqlalchemy<2.1,>=1
 Requires-Dist: verlib2==0.2
+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.37.0/src/sqlalchemy_cratedb.egg-info/SOURCES.txt

@@ -0,0 +1,36 @@
+CHANGES.md
+CODE_OF_CONDUCT.md
+CONTRIBUTING.md
+DEVELOP.md
+LICENSE
+MANIFEST.in
+NOTICE
+README.md
+backlog.md
+pyproject.toml
+setup.cfg
+docs/requirements.txt
+docs/_extra/robots.txt
+src/sqlalchemy_cratedb/__init__.py
+src/sqlalchemy_cratedb/compiler.py
+src/sqlalchemy_cratedb/dialect.py
+src/sqlalchemy_cratedb/predicate.py
+src/sqlalchemy_cratedb/sa_version.py
+src/sqlalchemy_cratedb/support.py
+src/sqlalchemy_cratedb.egg-info/PKG-INFO
+src/sqlalchemy_cratedb.egg-info/SOURCES.txt
+src/sqlalchemy_cratedb.egg-info/dependency_links.txt
+src/sqlalchemy_cratedb.egg-info/entry_points.txt
+src/sqlalchemy_cratedb.egg-info/requires.txt
+src/sqlalchemy_cratedb.egg-info/top_level.txt
+src/sqlalchemy_cratedb/compat/__init__.py
+src/sqlalchemy_cratedb/compat/api13.py
+src/sqlalchemy_cratedb/compat/core10.py
+src/sqlalchemy_cratedb/compat/core14.py
+src/sqlalchemy_cratedb/compat/core20.py
+src/sqlalchemy_cratedb/type/__init__.py
+src/sqlalchemy_cratedb/type/array.py
+src/sqlalchemy_cratedb/type/geo.py
+src/sqlalchemy_cratedb/type/object.py
+src/sqlalchemy_cratedb/type/vector.py
+tests/test_schema.py

{sqlalchemy_cratedb-0.36.0 → sqlalchemy_cratedb-0.37.0}/src/sqlalchemy_cratedb.egg-info/requires.txt

@@ -1,4 +1,4 @@
-crate
+crate==1.0.0dev0
 geojson<4,>=2.5
 sqlalchemy<2.1,>=1
 verlib2==0.2
@@ -7,6 +7,9 @@ verlib2==0.2
 backports.zoneinfo<1
 importlib-resources
 
+[all]
+sqlalchemy-cratedb[vector]
+
 [develop]
 black<25
 mypy<1.11
@@ -24,9 +27,13 @@ build<2
 twine<6
 
 [test]
+cratedb-toolkit[testing]
 dask[dataframe]
 pandas<2.3
 pueblo>=0.0.7
 pytest<9
 pytest-cov<6
 pytest-mock<4
+
+[vector]
+numpy

sqlalchemy_cratedb-0.37.0/tests/test_schema.py

@@ -0,0 +1,25 @@
+import sqlalchemy as sa
+
+from tests.conftest import TESTDRIVE_DATA_SCHEMA
+
+
+def test_correct_schema(cratedb_service):
+    """
+    Tests that the correct schema is being picked up.
+    """
+    database = cratedb_service.database
+
+    tablename = f'"{TESTDRIVE_DATA_SCHEMA}"."foobar"'
+    inspector: sa.Inspector = sa.inspect(database.engine)
+    database.run_sql(f"CREATE TABLE {tablename} AS SELECT 1")
+
+    assert TESTDRIVE_DATA_SCHEMA in inspector.get_schema_names()
+
+    table_names = inspector.get_table_names(schema=TESTDRIVE_DATA_SCHEMA)
+    assert table_names == ["foobar"]
+
+    view_names = inspector.get_view_names(schema=TESTDRIVE_DATA_SCHEMA)
+    assert view_names == []
+
+    indexes = inspector.get_indexes(tablename)
+    assert indexes == []