diff --git a/Cargo.lock b/Cargo.lock index 0f61afe9..d47615e2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,8 +11,7 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "aerospike-core" version = "3.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b902ab303a25e4df94d84089d40fb0e4590eba57546831a7a63c98b68fcedb8" +source = "git+https://github.com/aerospike/aerospike-client-rust?rev=51d5a36#51d5a36d5d1b65c27e0688344bcce110910c90c5" dependencies = [ "aerospike-macro", "aerospike-rt", @@ -40,8 +39,7 @@ dependencies = [ [[package]] name = "aerospike-macro" version = "3.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f949bc72ae7e37122740c02bf5f33f9e1e1512b3dcf2859de2d8dd6e5433376e" +source = "git+https://github.com/aerospike/aerospike-client-rust?rev=51d5a36#51d5a36d5d1b65c27e0688344bcce110910c90c5" dependencies = [ "proc-macro2", "quote", @@ -51,8 +49,7 @@ dependencies = [ [[package]] name = "aerospike-rt" version = "3.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "150a667e7595c8d0d0845045771e97fcb089c60521f091d4b3dd33982283bcb1" +source = "git+https://github.com/aerospike/aerospike-client-rust?rev=51d5a36#51d5a36d5d1b65c27e0688344bcce110910c90c5" dependencies = [ "tokio", ] @@ -63,11 +60,12 @@ version = "1.0.0" dependencies = [ "aerospike-core", "futures", + "half", "hex", "indexmap", "log", "mimalloc", - "ordered-float 3.9.2", + "ordered-float", "parking_lot", "pyo3", "pyo3-async-runtimes", @@ -502,6 +500,17 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -681,16 +690,6 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3e2e65a1a2e43cfcb47a895c4c8b10d1f4a61097f9f254f183aee60cad9c651d" -[[package]] -name = "matrixmultiply" -version = "0.3.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" -dependencies = [ - "autocfg", - "rawpointer", -] - [[package]] name = "memchr" version = "2.8.2" @@ -727,21 +726,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "ndarray" -version = "0.17.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" -dependencies = [ - "matrixmultiply", - "num-complex", - "num-integer", - "num-traits", - "portable-atomic", - "portable-atomic-util", - "rawpointer", -] - [[package]] name = "num-bigint" version = "0.4.6" @@ -785,22 +769,6 @@ dependencies = [ "autocfg", ] -[[package]] -name = "numpy" -version = "0.29.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4" -dependencies = [ - "libc", - "ndarray", - "num-complex", - "num-integer", - "num-traits", - "pyo3", - "pyo3-build-config", - "rustc-hash 2.1.2", -] - [[package]] name = "once_cell" version = "1.21.4" @@ -816,15 +784,6 @@ dependencies = [ "num-traits", ] -[[package]] -name = "ordered-float" -version = "5.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7d950ca161dc355eaf28f82b11345ed76c6e1f6eb1f4f4479e0323b9e2fbd0e" -dependencies = [ - "num-traits", -] - [[package]] name = "parking" version = "2.2.1" @@ -904,15 +863,6 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" -[[package]] -name = "portable-atomic-util" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" -dependencies = [ - "portable-atomic", -] - [[package]] name = "powerfmt" version = "0.2.0" @@ -1028,15 +978,12 @@ checksum = "a267ea4da9a831f534def7f1d8e14f581d5640d3e5af6527072724f216f4dbbf" dependencies = [ "anyhow", "chrono", - "either", "indexmap", "inventory", "itertools 0.14.0", "log", "maplit", "num-complex", - "numpy", - "ordered-float 5.3.0", "pyo3", "pyo3-stub-gen-derive", "rustpython-parser", @@ -1122,12 +1069,6 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" -[[package]] -name = "rawpointer" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" - [[package]] name = "redox_syscall" version = "0.5.18" @@ -1195,12 +1136,6 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" -[[package]] -name = "rustc-hash" -version = "2.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" - [[package]] name = "rustls" version = "0.23.45" @@ -1263,7 +1198,7 @@ dependencies = [ "num-traits", "phf", "phf_codegen", - "rustc-hash 1.1.0", + "rustc-hash", "rustpython-ast", "rustpython-parser-core", "tiny-keccak", diff --git a/Cargo.toml b/Cargo.toml index a1899b4b..ba86a3d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ crate-type = ["cdylib", "rlib"] pyo3 = { version = "0.29", default-features = false, features = ["indexmap"] } pyo3-async-runtimes = { version = "0.29", features = ["tokio-runtime"] } tokio = "1.40" -aerospike-core = { version = "3.0.0", features = ["rt-tokio", "dynamic-config", "serialization"] } +aerospike-core = { git = "https://github.com/aerospike/aerospike-client-rust", rev = "51d5a36", package = "aerospike-core", features = ["rt-tokio", "dynamic-config", "serialization"] } hex = "0.4" # Metrics snapshots re-serialize through core's `serialization` feature; these # convert the serde tree into plain Python dicts for `to_dict()`. @@ -24,7 +24,7 @@ parking_lot = "0.12" # match aerospike-core's PartitionStatus mutex (v3) # ordered Python dicts. indexmap = "2" ordered-float = { version = "3", default-features = false } -pyo3-stub-gen = "0.23" +pyo3-stub-gen = { version = "0.23", default-features = false, features = ["infer_signature"] } log = "0.4" pyo3-log = "0.13" futures = "0.3" @@ -37,6 +37,8 @@ webpki-roots = { version = "1.0", optional = true } # exits, and on macOS shares a fixed TLS slot with any other bundled v3. mimalloc = { version = "0.1", features = ["v2"] } rand = { version = "0.8", features = ["small_rng"] } +# Converts FLOAT16 vector elements to and from Python floats. +half = "2" [build-dependencies] # Re-exports pyo3's interpreter cfgs (Py_3_13, Py_GIL_DISABLED, …) via build.rs diff --git a/README.md b/README.md index 127a255e..421d2f34 100644 --- a/README.md +++ b/README.md @@ -243,6 +243,53 @@ policy = WritePolicy() policy.use_compression = True ``` +## Vectors and Top-K queries + +Aerospike 8.2.1+ stores dense numeric vectors as bin values and ranks query +results server-side. Every cluster node must run 8.2.1 or later; otherwise the +client refuses the command before sending it. + +A `Vector` holds `FLOAT16`, `INT32`, `FLOAT32` (the default), or `FLOAT64` +elements, built from a sequence or from a 1-D buffer (a numpy array, +`array.array` or `memoryview`) whose format selects the element type. It is an +immutable sequence, and `numpy.asarray(vector)` reads its elements without +copying. It is stored as a whole bin value, not nested in a list or map. + +```python +from aerospike_native import ( + BinDataType, ExpOperation, ExpReadFlags, FilterExpression, Operation, Order, + PartitionFilter, Statement, Vector, VectorDistanceMetric, +) + +await client.put(key, {"name": "lamp", "embedding": Vector([0.12, 0.98, -0.34])}) +``` + +A Top-K query keeps the `k` best records across the whole cluster, ordered by +up to two integer, double or bytes keys. Order by a vector distance by computing it in the query's +operations and ordering on the computed bin: + +```python +query = Vector([0.10, 0.95, -0.30]) +distance = FilterExpression.vector_distance( + VectorDistanceMetric.COSINE_SIMILARITY, query, FilterExpression.vector_bin("embedding"), +) + +stmt = Statement("test", "products") +stmt.set_operations([ + Operation.get_bin("name"), + ExpOperation.read("similarity", distance, ExpReadFlags.DEFAULT), +]) +stmt.set_order_by("similarity", BinDataType.DOUBLE, Order.DESC) +stmt.top_k = 10 + +async for record in await client.query(stmt, PartitionFilter.all()): + print(record.bins["name"], record.bins["similarity"]) +``` + +Records are delivered once every node has answered. Top-K runs only through +`query` (not background queries), and is incompatible with `max_records`, +`include_bin_data=False`, and a resumed partition filter. + ## Versioning The PNC follows [SemVer](https://semver.org/). Pre-releases use the diff --git a/python/aerospike_native/_native.pyi b/python/aerospike_native/_native.pyi index d2d92104..dcd3ca60 100644 --- a/python/aerospike_native/_native.pyi +++ b/python/aerospike_native/_native.pyi @@ -3,6 +3,7 @@ import _native import builtins +import collections.abc import enum import typing __all__ = [ @@ -21,6 +22,7 @@ __all__ = [ "BatchUDFPolicy", "BatchWriteOp", "BatchWritePolicy", + "BinDataType", "BitOperation", "BitPolicy", "BitWriteFlags", @@ -83,6 +85,8 @@ __all__ = [ "Node", "NodeMetricsSnapshot", "Operation", + "Order", + "OrderBy", "PartitionFilter", "PartitionStatus", "Privilege", @@ -123,6 +127,9 @@ __all__ = [ "UDFLang", "UdfRemoveTask", "User", + "Vector", + "VectorDistanceMetric", + "VectorElementType", "Version", "WritePolicy", "_LocalClient", @@ -2617,6 +2624,35 @@ class FilterExpression: Create a HLL bin expression """ @staticmethod + def vector_bin(name: builtins.str) -> _native.FilterExpression: + r""" + Create a vector bin expression, the stored side of :meth:`vector_distance`. + """ + @staticmethod + def vector_distance(metric: _native.VectorDistanceMetric, query: _native.Vector, bin: _native.FilterExpression) -> _native.FilterExpression: + r""" + Create an expression that returns, as a float, the distance between + ``query`` and a stored vector under ``metric``. + + ``bin`` is typically :meth:`vector_bin`. A stored value that is not a + vector, or whose element type or dimension count differs from + ``query``, evaluates to unknown: the record does not match as a filter, + and an ``ExpOperation.read`` with ``ExpReadFlags.EVAL_NO_FAIL`` returns + no value for it. Requires every node to run Aerospike Server 8.2.1 or + later. + + Example:: + + close = FilterExpression.lt( + FilterExpression.vector_distance( + VectorDistanceMetric.EUCLIDEAN_SQUARED, + Vector([0.1, 0.9, 0.3]), + FilterExpression.vector_bin("embedding"), + ), + FilterExpression.float_val(0.25), + ) + """ + @staticmethod def bin_exists(name: builtins.str) -> _native.FilterExpression: r""" Create function that returns if bin of specified name exists. @@ -5391,6 +5427,31 @@ class Operation: Create a Prepend operation (byte-level prepend to a string or bytes bin). """ +@typing.final +class OrderBy: + r""" + One Top-K order key of a :class:`Statement`, set with + :meth:`Statement.set_order_by` or :meth:`Statement.then_order_by`. + """ + @property + def bin_name(self) -> builtins.str: + r""" + Name of the returned bin the key reads. + """ + @property + def data_type(self) -> _native.BinDataType: + r""" + Declared type of the key. + """ + @property + def order(self) -> _native.Order: + r""" + Sort direction. + """ + def __eq__(self, other: typing.Any) -> builtins.bool: ... + def __hash__(self) -> builtins.int: ... + def __repr__(self) -> builtins.str: ... + @typing.final class PartitionFilter: @property @@ -6219,6 +6280,20 @@ class Statement: def set_name(self) -> typing.Optional[builtins.str]: ... @set_name.setter def set_name(self, value: typing.Optional[builtins.str]) -> None: ... + @property + def order_by(self) -> builtins.list[_native.OrderBy]: + r""" + The Top-K order keys, in comparison order; empty for an unranked query. + """ + @property + def top_k(self) -> typing.Optional[builtins.int]: + r""" + The Top-K limit ``k``, in ``[1, 1000]``, or ``None`` for an unranked + query. Requires an order key (see :meth:`set_order_by`); the two may be + set in either order. + """ + @top_k.setter + def top_k(self, value: typing.Optional[builtins.int]) -> None: ... def __new__(cls, namespace: builtins.str, set_name: typing.Optional[builtins.str] = None, bins: typing.Optional[typing.Sequence[builtins.str]] = None) -> _native.Statement: ... def set_aggregate_function(self, package_name: builtins.str, function_name: builtins.str, function_args: typing.Optional[typing.Sequence[typing.Any]] = None) -> None: r""" @@ -6240,6 +6315,36 @@ class Statement: 8.1.2 only accept the basic ``Read`` op here; 8.1.2+ also accepts CDT, expression, bit, and HLL reads. """ + def set_order_by(self, bin_name: builtins.str, data_type: _native.BinDataType, order: _native.Order) -> None: + r""" + Sets the order key of a Top-K query, replacing any keys set before. + + A Top-K query returns the ``k`` best records by this key (see + :attr:`top_k`), best first. Every node ranks its own records and the + client merges them; a record that two nodes both returned keeps its + later-updated copy. Ranking runs after the index filter and + ``QueryPolicy.filter_expression``, so it changes which matches come + back, not which records match. + + ``bin_name`` names a bin of the *returned* record: a stored bin named in + ``bins``, or one produced by exactly one :meth:`set_operations` read, + such as an ``ExpOperation.read`` of :meth:`FilterExpression.vector_distance`. + A record whose key is missing or of another type ranks last in either + direction; ties break by digest. + + Checked when the query runs, which raises :class:`ValueError` for a bin + name that is empty, longer than 15 UTF-8 bytes or holds a NUL; a key the + projection does not return; or an order key without :attr:`top_k`. + Top-K requires every node to run Aerospike Server 8.2.1 or later, and + cannot be combined with an aggregation, a background query, + ``QueryPolicy.max_records`` or a resumed partition filter. + """ + def then_order_by(self, bin_name: builtins.str, data_type: _native.BinDataType, order: _native.Order) -> None: + r""" + Adds a second order key, compared when the first key ties. It must name + a different bin; a Top-K query takes at most two keys. See + :meth:`set_order_by`. + """ class StringOperation: r""" @@ -6953,6 +7058,85 @@ class User: Number of currently open connections for the user """ +@typing.final +class Vector: + r""" + A dense numeric vector, stored in a bin as a ``VECTOR`` value and compared + with :meth:`FilterExpression.vector_distance`. + + Build one from a sequence of numbers, with an optional element type + (default :attr:`VectorElementType.FLOAT32`), or from a 1-D buffer such as a + numpy array, ``array.array`` or ``memoryview``, whose format selects the + element type:: + + embedding = Vector([0.12, 0.98, -0.34]) + counts = Vector([3, 0, 7], VectorElementType.INT32) + compact = Vector(np.array([0.5, 1.5], dtype=np.float16)) + + A vector has at least one element, and every floating-point element is + finite: an empty sequence, a NaN or an infinity raises :class:`ValueError`. + ``FLOAT16`` elements given as Python floats are rounded to the nearest + half-precision value. + + A vector is an immutable, hashable sequence: ``len(v)``, ``v[i]``, slicing + and iteration work as they do for a tuple, and two vectors are equal when + their element types and elements are equal. It also exposes its elements + as a read-only buffer, so ``numpy.asarray(v)`` and ``memoryview(v)`` read + them in place without copying. + + Writing a vector, or sending a vector expression, requires every cluster + node to run Aerospike Server 8.2.1 or later; otherwise the client refuses + the command before sending it. The server stores a vector only as a whole + bin value, not nested in a list or map bin. + """ + @property + def element_type(self) -> _native.VectorElementType: + r""" + The element type. + """ + @property + def dimensions(self) -> builtins.int: + r""" + The number of elements. + """ + @property + def value(self) -> builtins.list[builtins.int] | builtins.list[builtins.float]: + r""" + The elements as a new list of ``int`` (``INT32``) or ``float`` values. + """ + def __new__(cls, data: _native.Vector | typing.Sequence[builtins.float] | collections.abc.Buffer, element_type: typing.Optional[_native.VectorElementType] = None) -> _native.Vector: ... + @staticmethod + def from_wire_bytes(data: builtins.bytes) -> _native.Vector: + r""" + Decodes a complete vector wire value: a version byte, an element-type + byte, the dimension count as a little-endian ``u32``, then the + little-endian elements. Use it for a vector carried as ``bytes``, such + as one an application stored in a blob bin. + + Example:: + + record = await client.get(key, ["embedding_blob"]) + embedding = Vector.from_wire_bytes(record.bins["embedding_blob"]) + + Args: + data: The encoded vector. + + Returns: + The decoded vector. + + Raises: + ValueError: The bytes are not a valid vector wire value. + """ + def __len__(self) -> builtins.int: ... + def __getitem__(self, index: typing.Any) -> builtins.int | builtins.float | builtins.list[builtins.int] | builtins.list[builtins.float]: ... + def __iter__(self) -> typing.Iterator[builtins.int | builtins.float]: ... + def __eq__(self, other: object) -> builtins.bool: ... + def __ne__(self, other: object) -> builtins.bool: ... + def __hash__(self) -> builtins.int: ... + def __repr__(self) -> builtins.str: ... + def __buffer__(self, flags: int) -> memoryview: ... + def __release_buffer__(self, buffer: memoryview) -> None: ... + class Version: @property def major(self) -> builtins.int: ... @@ -6991,6 +7175,16 @@ class Version: Returns true if server supports two-phase server query selection (field ``44`` WHERE explain → execute). Requires server >= 8.2.0. """ + def supports_vector(self) -> builtins.bool: + r""" + Returns true if server supports ``VECTOR`` values and vector-distance + expressions. Requires server >= 8.2.1. + """ + def supports_query_top_k(self) -> builtins.bool: + r""" + Returns true if server supports Top-K queries (the order-by and + Top-K query fields). Requires server >= 8.2.1. + """ def __str__(self) -> builtins.str: ... def __repr__(self) -> builtins.str: ... @@ -7252,6 +7446,26 @@ class AuthMode(enum.Enum): Requires server version 5.7.0+ """ +@typing.final +class BinDataType(enum.Enum): + r""" + Declared type of a Top-K order key, used with :meth:`Statement.set_order_by`. + + A record whose order-key bin is missing or holds another type ranks last. + """ + INTEGER = ... + r""" + 64-bit signed integer. + """ + DOUBLE = ... + r""" + 64-bit float. NaN ranks after every other value; ``-0.0`` equals ``0.0``. + """ + BYTES = ... + r""" + Bytes, compared as unsigned bytes; a prefix ranks before the longer value. + """ + @typing.final class BitWriteFlags(enum.IntEnum): DEFAULT: builtins.int @@ -7390,6 +7604,7 @@ class ExpType(enum.Enum): FLOAT = ... GEO = ... HLL = ... + VECTOR = ... @typing.final class ExpWriteFlags(enum.IntEnum): @@ -7601,6 +7816,21 @@ class MapWriteMode(enum.Enum): If the key does not exist, a new item will be created. """ +@typing.final +class Order(enum.Enum): + r""" + Sort direction of a Top-K order key: which ``k`` records are kept and the + order they are returned in. + """ + ASC = ... + r""" + Smallest values first. + """ + DESC = ... + r""" + Largest values first. + """ + @typing.final class PrivilegeCode(enum.Enum): r""" @@ -7905,6 +8135,46 @@ class UDFLang(enum.Enum): Lua embedded programming language. """ +@typing.final +class VectorDistanceMetric(enum.Enum): + r""" + Metric of :meth:`FilterExpression.vector_distance`. + """ + EUCLIDEAN_SQUARED = ... + r""" + Squared Euclidean distance. Smaller is closer. + """ + DOT_PRODUCT = ... + r""" + Dot product. Larger is more similar. + """ + COSINE_SIMILARITY = ... + r""" + Cosine similarity. Larger is more similar. + """ + +@typing.final +class VectorElementType(enum.Enum): + r""" + Element type of a :class:`Vector`: how each element is stored on the wire. + """ + FLOAT16 = ... + r""" + IEEE 754 half precision. + """ + INT32 = ... + r""" + 32-bit signed integer. + """ + FLOAT32 = ... + r""" + 32-bit IEEE 754 float. + """ + FLOAT64 = ... + r""" + 64-bit IEEE 754 double. + """ + def geojson(geo_str: builtins.str) -> _native.GeoJSON: r""" Convert a GeoJSON string or coordinate pair to a GeoJSON object. diff --git a/python/conftest.py b/python/conftest.py index 260dc600..7d2988ef 100644 --- a/python/conftest.py +++ b/python/conftest.py @@ -251,6 +251,34 @@ async def supports_query_selection(aerospike_host, use_services_alternate): ) +@pytest_asyncio.fixture(scope="session", loop_scope="session") +async def supports_vector(aerospike_host, use_services_alternate): + """``True`` when every connected node accepts vector values and expressions. + + Probes ``Node.version.supports_vector()`` on all nodes (server >= 8.2.1). + The client refuses vector commands to a cluster where any node lacks it. + """ + return await _probe_all_nodes_version_capability( + aerospike_host, + use_services_alternate, + lambda version: version.supports_vector(), + ) + + +@pytest_asyncio.fixture(scope="session", loop_scope="session") +async def supports_query_top_k(aerospike_host, use_services_alternate): + """``True`` when every connected node supports Top-K queries. + + Probes ``Node.version.supports_query_top_k()`` on all nodes (server >= 8.2.1). + The client refuses a Top-K query to a cluster where any node lacks it. + """ + return await _probe_all_nodes_version_capability( + aerospike_host, + use_services_alternate, + lambda version: version.supports_query_top_k(), + ) + + @pytest.fixture(scope="session") def supports_query_selection_sync(aerospike_host, use_services_alternate): """Sync session gate for query selection (blocking integration tests).""" diff --git a/python/tests/integration/get_node_test.py b/python/tests/integration/get_node_test.py index f7c67229..65a2208f 100644 --- a/python/tests/integration/get_node_test.py +++ b/python/tests/integration/get_node_test.py @@ -182,6 +182,13 @@ async def test_version_supports_query_selection(self, client): if (version.major, version.minor, version.patch) >= (8, 2, 0): assert result is True + async def test_version_supports_vector_and_top_k(self, client): + """Test Version.supports_vector() and supports_query_top_k() (server 8.2.1+).""" + version = client.nodes()[0].version + expected = (version.major, version.minor, version.patch) >= (8, 2, 1) + assert version.supports_vector() is expected + assert version.supports_query_top_k() is expected + async def test_version_supports_server_compiled_ael(self, client): """Test Version.supports_server_compiled_ael() method.""" nodes = client.nodes() diff --git a/python/tests/integration/mrt_test.py b/python/tests/integration/mrt_test.py index 12387757..987e3d21 100644 --- a/python/tests/integration/mrt_test.py +++ b/python/tests/integration/mrt_test.py @@ -51,6 +51,7 @@ ResultCode, Txn, TxnState, + UDFLang, WritePolicy, new_client, ) @@ -189,6 +190,22 @@ async def test_txn_abort_rolls_back(sc_client, sc_key): assert await _get_bin(sc_client, sc_key, "bin") == 1 +async def test_txn_abort_rolls_back_a_udf_write(sc_client, sc_key): + udf_path = os.path.join(os.path.dirname(__file__), "udf", "record_example.lua") + task = await sc_client.register_udf_from_file(udf_path, "record_example.lua", UDFLang.LUA) + assert await task.wait_till_complete() + await sc_client.put(sc_key, {"bin": 1}, policy=WritePolicy()) + + txn = Txn() + wp = WritePolicy() + wp.txn = txn + await sc_client.execute_udf(sc_key, "record_example", "writeBin", ["bin", 2], policy=wp) + + status = await sc_client.abort(txn) + assert status == AbortStatus.OK + assert await _get_bin(sc_client, sc_key, "bin") == 1 + + # --------------------------------------------------------------------------- # State transitions # --------------------------------------------------------------------------- diff --git a/python/tests/integration/udf/vector_example.lua b/python/tests/integration/udf/vector_example.lua new file mode 100644 index 00000000..67c5eb64 --- /dev/null +++ b/python/tests/integration/udf/vector_example.lua @@ -0,0 +1,17 @@ +-- Return the argument unchanged. +function echo(r, value) + return value +end + +-- Return a stored bin. +function read_bin(r, name) + return r[name] +end + +-- Stream a stored bin of each record. +function stream_bin(stream, name) + local function mapper(rec) + return rec[name] + end + return stream : map(mapper) +end diff --git a/python/tests/integration/vector_test.py b/python/tests/integration/vector_test.py new file mode 100644 index 00000000..aa6f1c15 --- /dev/null +++ b/python/tests/integration/vector_test.py @@ -0,0 +1,315 @@ +# Copyright 2023-2026 Aerospike, Inc. +# +# Portions may be licensed to Aerospike, Inc. under one or more contributor +# license agreements WHICH ARE COMPATIBLE WITH THE APACHE LICENSE, VERSION 2.0. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +"""Vector bins, vector-distance expressions, and Top-K queries.""" + +import array +import os + +import pytest +from aerospike_native import ( + BinDataType, + ExpOperation, + ExpReadFlags, + FilterExpression, + Key, + ListOperation, + ListPolicy, + Operation, + Order, + PartitionFilter, + QueryPolicy, + ResultCode, + Statement, + UDFLang, + Vector, + VectorDistanceMetric, + VectorElementType, +) +from aerospike_native.exceptions import InvalidRequest, ServerError +from aerospike_native.exceptions import ValueError as AerospikeValueError +from fixtures import TestFixtureCleanDB + +NAMESPACE = "test" +SET = "test" + + +def distance(metric, query, bin_name="embedding"): + return FilterExpression.vector_distance(metric, query, FilterExpression.vector_bin(bin_name)) + + +async def collect(client, stmt, policy=None): + records = await client.query(stmt, PartitionFilter.all(), policy=policy or QueryPolicy()) + return [record async for record in records] + + +class TestVectorBins(TestFixtureCleanDB): + @pytest.fixture(autouse=True) + def require_vector(self, supports_vector): + if not supports_vector: + pytest.skip("cluster lacks vector support (Node.version.supports_vector() is False)") + + @pytest.mark.parametrize( + "vector", + [ + Vector([1, -2, 2**31 - 1], VectorElementType.INT32), + Vector([0.5, -1.25, 3.0], VectorElementType.FLOAT32), + Vector([0.1, -2.5, 1e300], VectorElementType.FLOAT64), + Vector([0.5, -2.0, 65504.0], VectorElementType.FLOAT16), + ], + ids=["int32", "float32", "float64", "float16"], + ) + async def test_put_get_round_trip(self, client, key, vector): + await client.put(key, {"embedding": vector}) + record = await client.get(key) + assert record.bins["embedding"] == vector + assert record.bins["embedding"].element_type == vector.element_type + + async def test_buffer_round_trip(self, client, key): + embedding = array.array("f", (i / 768 for i in range(768))) + await client.put(key, {"embedding": Vector(embedding)}) + record = await client.get(key) + assert memoryview(record.bins["embedding"]) == memoryview(embedding) + + async def test_overwrite_with_different_dimensions(self, client, key): + await client.put(key, {"embedding": Vector([1.0, 2.0, 3.0])}) + replacement = Vector([4.0, 5.0], VectorElementType.FLOAT64) + await client.put(key, {"embedding": replacement}) + assert (await client.get(key)).bins["embedding"] == replacement + + async def test_operate_put_and_get(self, client, key): + vector = Vector([7, 8, 9], VectorElementType.INT32) + await client.operate(key, [Operation.put("embedding", vector)]) + record = await client.operate(key, [Operation.get_bin("embedding")]) + assert record.bins["embedding"] == vector + + async def test_batch_operate_reads_each_vector(self, client): + keys = [Key(NAMESPACE, SET, i) for i in range(3)] + vectors = [Vector([float(i), float(i + 1)]) for i in range(3)] + writes = [[Operation.put("embedding", v)] for v in vectors] + for result in await client.batch_operate(keys, writes, batch_policy=None, write_policy=None): + assert result.result_code == ResultCode.OK + + reads = await client.batch_operate( + keys, [[Operation.get_bin("embedding")]] * len(keys), batch_policy=None, write_policy=None, + ) + assert [result.record.bins["embedding"] for result in reads] == vectors + + @pytest.mark.parametrize( + "bins", + [{"items": [Vector([1.0])]}, {"lookup": {"k": Vector([1.0])}}], + ids=["list", "map"], + ) + async def test_vector_nested_in_collection_is_rejected_by_server(self, client, key, bins): + with pytest.raises(InvalidRequest) as exc_info: + await client.put(key, bins) + assert exc_info.value.result_code == ResultCode.PARAMETER_ERROR + + async def test_vector_as_list_append_value_is_rejected_by_server(self, client, key): + with pytest.raises(InvalidRequest) as exc_info: + await client.operate(key, [ListOperation.append("items", Vector([1.0]), ListPolicy())]) + assert exc_info.value.result_code == ResultCode.PARAMETER_ERROR + + @pytest.fixture + async def vector_udf(self, client): + udf_path = os.path.join(os.path.dirname(__file__), "udf", "vector_example.lua") + task = await client.register_udf_from_file(udf_path, "vector_example.lua", UDFLang.LUA) + assert await task.wait_till_complete() + yield "vector_example" + remove_task = await client.remove_udf("vector_example.lua") + await remove_task.wait_till_complete() + + async def test_udf_returns_vector(self, client, key, vector_udf): + vector = Vector([1.5, -2.0], VectorElementType.FLOAT64) + await client.put(key, {"embedding": vector}) + assert await client.execute_udf(key, vector_udf, "echo", [vector]) == vector + assert await client.execute_udf(key, vector_udf, "read_bin", ["embedding"]) == vector + + async def test_aggregation_streams_vectors(self, client, vector_udf): + vectors = [Vector([float(i), -float(i)], VectorElementType.FLOAT64) for i in range(3)] + for i, vector in enumerate(vectors): + await client.put(Key(NAMESPACE, SET, i), {"embedding": vector}) + stmt = Statement(NAMESPACE, SET) + stmt.set_aggregate_function(vector_udf, "stream_bin", ["embedding"]) + streamed = [record.bins["SUCCESS"] for record in await collect(client, stmt)] + assert sorted(streamed, key=lambda v: v[0]) == vectors + + +class TestVectorDistance(TestFixtureCleanDB): + @pytest.fixture(autouse=True) + def require_vector(self, supports_vector): + if not supports_vector: + pytest.skip("cluster lacks vector support (Node.version.supports_vector() is False)") + + @pytest.fixture + async def stored(self, client, key): + await client.put(key, {"embedding": Vector([2.0, 0.0, 1.0], VectorElementType.FLOAT64)}) + return key + + @pytest.mark.parametrize( + "metric, expected", + [ + (VectorDistanceMetric.EUCLIDEAN_SQUARED, 9.0), + (VectorDistanceMetric.DOT_PRODUCT, 5.0), + (VectorDistanceMetric.COSINE_SIMILARITY, 5.0 / 70**0.5), + ], + ) + async def test_metric_value(self, client, stored, metric, expected): + query = Vector([1.0, 2.0, 3.0], VectorElementType.FLOAT64) + record = await client.operate(stored, [ExpOperation.read("d", distance(metric, query), ExpReadFlags.DEFAULT)]) + assert record.bins["d"] == pytest.approx(expected) + + async def test_filter_by_distance(self, client, stored): + query = Vector([1.0, 2.0, 3.0], VectorElementType.FLOAT64) + near = FilterExpression.lt( + distance(VectorDistanceMetric.EUCLIDEAN_SQUARED, query), FilterExpression.float_val(10.0) + ) + assert (await client.get(stored, filter_expression=near)).bins["embedding"] is not None + + far = FilterExpression.lt( + distance(VectorDistanceMetric.EUCLIDEAN_SQUARED, query), FilterExpression.float_val(1.0) + ) + with pytest.raises(ServerError) as exc_info: + await client.get(stored, filter_expression=far) + assert exc_info.value.result_code == ResultCode.FILTERED_OUT + + async def test_mismatched_dimensions_evaluate_to_unknown(self, client, stored): + mismatched = distance(VectorDistanceMetric.DOT_PRODUCT, Vector([1.0, 2.0], VectorElementType.FLOAT64)) + record = await client.operate(stored, [ExpOperation.read("d", mismatched, ExpReadFlags.EVAL_NO_FAIL)]) + assert record.bins.get("d") is None + + +class TestTopK(TestFixtureCleanDB): + @pytest.fixture(autouse=True) + def require_top_k(self, supports_query_top_k): + if not supports_query_top_k: + pytest.skip("cluster lacks Top-K (Node.version.supports_query_top_k() is False)") + + @pytest.fixture + async def products(self, client): + rows = [ + (1, 9.5, b"\x09", 40, [0.0, 0.0]), + (2, 7.25, b"\x07\x01", 90, [3.0, 4.0]), + (3, 7.0, b"\x07", 90, [1.0, 1.0]), + (4, 5.0, b"\x05", 70, [0.5, 0.2]), + (5, 3.0, b"\x03", None, [5.0, 5.0]), + ] + for pid, price, sku, score, embedding in rows: + bins = {"id": pid, "price": price, "sku": sku, + "embedding": Vector(embedding, VectorElementType.FLOAT64)} + if score is not None: + bins["score"] = score + await client.put(Key(NAMESPACE, SET, pid), bins) + return rows + + async def test_highest_scores_first(self, client, products): + stmt = Statement(NAMESPACE, SET, ["id", "score"]) + stmt.set_order_by("score", BinDataType.INTEGER, Order.DESC) + stmt.top_k = 3 + records = await collect(client, stmt) + assert [r.bins["score"] for r in records] == [90, 90, 70] + + async def test_second_key_breaks_ties(self, client, products): + stmt = Statement(NAMESPACE, SET, ["id", "price", "score"]) + stmt.set_order_by("score", BinDataType.INTEGER, Order.DESC) + stmt.then_order_by("price", BinDataType.DOUBLE, Order.ASC) + stmt.top_k = 2 + assert [r.bins["id"] for r in await collect(client, stmt)] == [3, 2] + + async def test_bytes_order_key(self, client, products): + stmt = Statement(NAMESPACE, SET, ["id", "sku"]) + stmt.set_order_by("sku", BinDataType.BYTES, Order.DESC) + stmt.top_k = 3 + # Descending, so the longer value b"\x07\x01" ranks before its prefix b"\x07". + assert [r.bins["id"] for r in await collect(client, stmt)] == [1, 2, 3] + + async def test_missing_key_ranks_last_and_k_beyond_matches_returns_all(self, client, products): + stmt = Statement(NAMESPACE, SET, ["id", "score"]) + stmt.set_order_by("score", BinDataType.INTEGER, Order.ASC) + stmt.top_k = 100 + ids = [r.bins["id"] for r in await collect(client, stmt)] + # Ids 2 and 3 tie on score, and a tie with no second key breaks by digest. + assert ids[:2] == [1, 4] + assert set(ids[2:4]) == {2, 3} + assert ids[4:] == [5] + + async def test_nearest_by_distance(self, client, products): + query = Vector([0.0, 0.0], VectorElementType.FLOAT64) + stmt = Statement(NAMESPACE, SET) + stmt.set_operations([ + Operation.get_bin("id"), + ExpOperation.read("d", distance(VectorDistanceMetric.EUCLIDEAN_SQUARED, query), ExpReadFlags.DEFAULT), + ]) + stmt.set_order_by("d", BinDataType.DOUBLE, Order.ASC) + stmt.top_k = 3 + records = await collect(client, stmt) + assert [r.bins["id"] for r in records] == [1, 4, 3] + assert [r.bins["d"] for r in records] == pytest.approx([0.0, 0.29, 2.0]) + + @pytest.mark.parametrize( + "configure", + [ + lambda stmt, policy: setattr(stmt, "top_k", 0), + lambda stmt, policy: setattr(stmt, "top_k", 1001), + lambda stmt, policy: setattr(policy, "max_records", 10), + lambda stmt, policy: setattr(policy, "include_bin_data", False), + lambda stmt, policy: stmt.then_order_by("score", BinDataType.INTEGER, Order.ASC), + ], + ids=["k-zero", "k-too-large", "max-records", "no-bin-data", "same-bin-twice"], + ) + async def test_invalid_top_k_is_rejected_before_sending(self, client, configure): + stmt = Statement(NAMESPACE, SET, ["score"]) + stmt.set_order_by("score", BinDataType.INTEGER, Order.DESC) + stmt.top_k = 5 + policy = QueryPolicy() + configure(stmt, policy) + with pytest.raises(AerospikeValueError) as exc_info: + await collect(client, stmt, policy) + assert exc_info.value.result_code == ResultCode.PARAMETER_ERROR + + async def test_background_query_rejects_top_k(self, client): + stmt = Statement(NAMESPACE, SET) + stmt.set_order_by("score", BinDataType.INTEGER, Order.DESC) + stmt.top_k = 5 + with pytest.raises(AerospikeValueError) as exc_info: + await client.query_operate(stmt, [Operation.put("seen", 1)]) + assert exc_info.value.result_code == ResultCode.PARAMETER_ERROR + + +class TestVectorUnsupportedCluster(TestFixtureCleanDB): + """A cluster with a node below 8.2.1 gets a client-side refusal, not a send.""" + + @pytest.fixture(autouse=True) + def require_pre_vector_cluster(self, supports_vector): + if supports_vector: + pytest.skip("every node supports vectors") + + async def test_vector_write_is_refused(self, client, key): + with pytest.raises(AerospikeValueError, match="not supported") as exc_info: + await client.put(key, {"embedding": Vector([1.0])}) + assert exc_info.value.result_code == ResultCode.PARAMETER_ERROR + + async def test_vector_expression_is_refused(self, client, key): + await client.put(key, {"n": 1}) + query = distance(VectorDistanceMetric.DOT_PRODUCT, Vector([1.0])) + with pytest.raises(AerospikeValueError, match="not supported"): + await client.operate(key, [ExpOperation.read("d", query, ExpReadFlags.DEFAULT)]) + + async def test_top_k_query_is_refused(self, client): + stmt = Statement(NAMESPACE, SET, ["n"]) + stmt.set_order_by("n", BinDataType.INTEGER, Order.DESC) + stmt.top_k = 5 + with pytest.raises(AerospikeValueError, match="Top-K is not supported"): + await collect(client, stmt) diff --git a/python/tests/unit/vector_test.py b/python/tests/unit/vector_test.py new file mode 100644 index 00000000..d18ba860 --- /dev/null +++ b/python/tests/unit/vector_test.py @@ -0,0 +1,295 @@ +# Copyright 2023-2026 Aerospike, Inc. +# +# Portions may be licensed to Aerospike, Inc. under one or more contributor +# license agreements WHICH ARE COMPATIBLE WITH THE APACHE LICENSE, VERSION 2.0. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +import array +import ctypes +import struct +import sys + +import pytest + +import aerospike_native +from aerospike_native import ( + BinDataType, + FilterExpression, + Order, + Statement, + Vector, + VectorDistanceMetric, + VectorElementType, +) + +NON_NATIVE_FLOAT = ctypes.c_float.__ctype_be__ if sys.byteorder == "little" else ctypes.c_float.__ctype_le__ + + +def wire_bytes(type_code, fmt, values, version=1): + """A vector wire value: version, element type, LE dimensions, LE elements.""" + return bytes([version, type_code]) + struct.pack(f"= "3.12" pyperf==2.9.0 pytest==9.1.1 pytest_asyncio==1.4.0 diff --git a/src/cluster.rs b/src/cluster.rs index 0936fa4f..70e98898 100644 --- a/src/cluster.rs +++ b/src/cluster.rs @@ -163,6 +163,18 @@ use crate::policies::AdminPolicy; self._as.supports_query_selection() } + /// Returns true if server supports ``VECTOR`` values and vector-distance + /// expressions. Requires server >= 8.2.1. + pub fn supports_vector(&self) -> bool { + self._as.supports_vector() + } + + /// Returns true if server supports Top-K queries (the order-by and + /// Top-K query fields). Requires server >= 8.2.1. + pub fn supports_query_top_k(&self) -> bool { + self._as.supports_query_top_k() + } + pub fn __str__(&self) -> String { format!("{}.{}.{}.{}", self._as.major, self._as.minor, self._as.patch, self._as.build) } diff --git a/src/enums.rs b/src/enums.rs index 9b10ac39..f979f60e 100644 --- a/src/enums.rs +++ b/src/enums.rs @@ -1609,3 +1609,109 @@ pub enum Concurrency { #[classattr] const NO_FAIL: i64 = 0x10; } + + //////////////////////////////////////////////////////////////////////////////////////////// + // + // Top-K: BinDataType, Order + // + //////////////////////////////////////////////////////////////////////////////////////////// + + /// Declared type of a Top-K order key, used with :meth:`Statement.set_order_by`. + /// + /// A record whose order-key bin is missing or holds another type ranks last. + #[gen_stub_pyclass_enum(module = "_native")] + #[pyclass(from_py_object, module = "_native")] + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + pub enum BinDataType { + /// 64-bit signed integer. + #[pyo3(name = "INTEGER")] + Integer, + /// 64-bit float. NaN ranks after every other value; ``-0.0`` equals ``0.0``. + #[pyo3(name = "DOUBLE")] + Double, + /// Bytes, compared as unsigned bytes; a prefix ranks before the longer value. + #[pyo3(name = "BYTES")] + Bytes, + } + + impl From for aerospike_core::OrderByType { + fn from(input: BinDataType) -> Self { + match input { + BinDataType::Integer => aerospike_core::OrderByType::Integer, + BinDataType::Double => aerospike_core::OrderByType::Double, + BinDataType::Bytes => aerospike_core::OrderByType::Bytes, + } + } + } + + impl From for BinDataType { + fn from(input: aerospike_core::OrderByType) -> Self { + match input { + aerospike_core::OrderByType::Integer => BinDataType::Integer, + aerospike_core::OrderByType::Double => BinDataType::Double, + aerospike_core::OrderByType::Bytes => BinDataType::Bytes, + aerospike_core::OrderByType::String => { + unreachable!("order keys are only set through BinDataType, which has no STRING") + } + } + } + } + + /// Sort direction of a Top-K order key: which ``k`` records are kept and the + /// order they are returned in. + #[gen_stub_pyclass_enum(module = "_native")] + #[pyclass(from_py_object, module = "_native")] + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + pub enum Order { + /// Smallest values first. + #[pyo3(name = "ASC")] + Asc, + /// Largest values first. + #[pyo3(name = "DESC")] + Desc, + } + + impl From for aerospike_core::Order { + fn from(input: Order) -> Self { + match input { + Order::Asc => aerospike_core::Order::Asc, + Order::Desc => aerospike_core::Order::Desc, + } + } + } + + impl From for Order { + fn from(input: aerospike_core::Order) -> Self { + match input { + aerospike_core::Order::Asc => Order::Asc, + aerospike_core::Order::Desc => Order::Desc, + } + } + } + + macro_rules! impl_simple_enum_eq_hash { + ($($ty:ty),*) => {$( + #[pymethods] + impl $ty { + fn __richcmp__(&self, other: &$ty, op: CompareOp) -> PyResult { + match op { + CompareOp::Eq => Ok(self == other), + CompareOp::Ne => Ok(self != other), + _ => Err(pyo3::exceptions::PyNotImplementedError::new_err( + "Only == and != comparisons are supported", + )), + } + } + + fn __hash__(&self) -> u64 { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut hasher = DefaultHasher::new(); + self.hash(&mut hasher); + hasher.finish() + } + } + )*}; + } + + impl_simple_enum_eq_hash!(BinDataType, Order); diff --git a/src/expressions.rs b/src/expressions.rs index b90a66b5..03a6a9cd 100644 --- a/src/expressions.rs +++ b/src/expressions.rs @@ -24,7 +24,7 @@ use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_py use crate::cdt::*; use crate::errors::RustClientError; -use crate::record::PythonValue; +use crate::record::{PythonValue, Vector}; use crate::string_ops::StringNumericType; //////////////////////////////////////////////////////////////////////////////////////////// @@ -58,6 +58,8 @@ use crate::string_ops::StringNumericType; Geo, #[pyo3(name = "HLL")] HLL, + #[pyo3(name = "VECTOR")] + Vector, } impl From<&ExpType> for aerospike_core::expressions::ExpType { @@ -73,10 +75,52 @@ use crate::string_ops::StringNumericType; ExpType::Float => aerospike_core::expressions::ExpType::Float, ExpType::Geo => aerospike_core::expressions::ExpType::Geo, ExpType::HLL => aerospike_core::expressions::ExpType::Hll, + ExpType::Vector => aerospike_core::expressions::ExpType::Vector, } } } + //////////////////////////////////////////////////////////////////////////////////////////// + // + // VectorDistanceMetric + // + //////////////////////////////////////////////////////////////////////////////////////////// + + /// Metric of :meth:`FilterExpression.vector_distance`. + #[gen_stub_pyclass_enum(module = "_native")] + #[pyclass(from_py_object, module = "_native")] + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + pub enum VectorDistanceMetric { + /// Squared Euclidean distance. Smaller is closer. + #[pyo3(name = "EUCLIDEAN_SQUARED")] + EuclideanSquared, + /// Dot product. Larger is more similar. + #[pyo3(name = "DOT_PRODUCT")] + DotProduct, + /// Cosine similarity. Larger is more similar. + #[pyo3(name = "COSINE_SIMILARITY")] + CosineSimilarity, + } + + #[pymethods] + impl VectorDistanceMetric { + fn __richcmp__(&self, other: &VectorDistanceMetric, op: pyo3::class::basic::CompareOp) -> PyResult { + match op { + pyo3::class::basic::CompareOp::Eq => Ok(self == other), + pyo3::class::basic::CompareOp::Ne => Ok(self != other), + _ => Err(pyo3::exceptions::PyNotImplementedError::new_err("Only == and != comparisons are supported")), + } + } + + fn __hash__(&self) -> u64 { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut hasher = DefaultHasher::new(); + self.hash(&mut hasher); + hasher.finish() + } + } + #[pymethods] impl ExpType { } @@ -516,6 +560,47 @@ use crate::string_ops::StringNumericType; } } + #[staticmethod] + /// Create a vector bin expression, the stored side of :meth:`vector_distance`. + pub fn vector_bin(name: String) -> Self { + FilterExpression { + _as: aerospike_core::expressions::vector_bin(name), + } + } + + #[staticmethod] + /// Create an expression that returns, as a float, the distance between + /// ``query`` and a stored vector under ``metric``. + /// + /// ``bin`` is typically :meth:`vector_bin`. A stored value that is not a + /// vector, or whose element type or dimension count differs from + /// ``query``, evaluates to unknown: the record does not match as a filter, + /// and an ``ExpOperation.read`` with ``ExpReadFlags.EVAL_NO_FAIL`` returns + /// no value for it. Requires every node to run Aerospike Server 8.2.1 or + /// later. + /// + /// Example:: + /// + /// close = FilterExpression.lt( + /// FilterExpression.vector_distance( + /// VectorDistanceMetric.EUCLIDEAN_SQUARED, + /// Vector([0.1, 0.9, 0.3]), + /// FilterExpression.vector_bin("embedding"), + /// ), + /// FilterExpression.float_val(0.25), + /// ) + pub fn vector_distance(metric: VectorDistanceMetric, query: PyRef<'_, Vector>, bin: FilterExpression) -> Self { + use aerospike_core::expressions::vector; + let build = match metric { + VectorDistanceMetric::EuclideanSquared => vector::euclidean_squared_distance, + VectorDistanceMetric::DotProduct => vector::dot_product, + VectorDistanceMetric::CosineSimilarity => vector::cosine_similarity, + }; + FilterExpression { + _as: build(&query.v, bin._as), + } + } + #[staticmethod] /// Create function that returns if bin of specified name exists. pub fn bin_exists(name: String) -> Self { diff --git a/src/filter.rs b/src/filter.rs index 1faa387d..b5c38380 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -28,7 +28,7 @@ use tokio::sync::Mutex; use aerospike_core::query::RecordStream; -use crate::CollectionIndexType; +use crate::{BinDataType, CollectionIndexType, Order}; use crate::cdt::CTX; use crate::errors::RustClientError; use crate::expressions::FilterExpression; @@ -460,6 +460,116 @@ use crate::record::{Key, PythonValue, Record}; self._as.set_operations(core_ops); Ok(()) } + + /// Sets the order key of a Top-K query, replacing any keys set before. + /// + /// A Top-K query returns the ``k`` best records by this key (see + /// :attr:`top_k`), best first. Every node ranks its own records and the + /// client merges them; a record that two nodes both returned keeps its + /// later-updated copy. Ranking runs after the index filter and + /// ``QueryPolicy.filter_expression``, so it changes which matches come + /// back, not which records match. + /// + /// ``bin_name`` names a bin of the *returned* record: a stored bin named in + /// ``bins``, or one produced by exactly one :meth:`set_operations` read, + /// such as an ``ExpOperation.read`` of :meth:`FilterExpression.vector_distance`. + /// A record whose key is missing or of another type ranks last in either + /// direction; ties break by digest. + /// + /// Checked when the query runs, which raises :class:`ValueError` for a bin + /// name that is empty, longer than 15 UTF-8 bytes or holds a NUL; a key the + /// projection does not return; or an order key without :attr:`top_k`. + /// Top-K requires every node to run Aerospike Server 8.2.1 or later, and + /// cannot be combined with an aggregation, a background query, + /// ``QueryPolicy.max_records`` or a resumed partition filter. + pub fn set_order_by(&mut self, bin_name: &str, data_type: BinDataType, order: Order) { + self._as.set_order_by(bin_name, data_type.into(), order.into()); + } + + /// Adds a second order key, compared when the first key ties. It must name + /// a different bin; a Top-K query takes at most two keys. See + /// :meth:`set_order_by`. + pub fn then_order_by(&mut self, bin_name: &str, data_type: BinDataType, order: Order) { + self._as.then_order_by(bin_name, data_type.into(), order.into()); + } + + /// The Top-K order keys, in comparison order; empty for an unranked query. + #[getter] + pub fn get_order_by(&self) -> Vec { + self._as.order_by.iter().map(|key| OrderBy { _as: key.clone() }).collect() + } + + /// The Top-K limit ``k``, in ``[1, 1000]``, or ``None`` for an unranked + /// query. Requires an order key (see :meth:`set_order_by`); the two may be + /// set in either order. + #[getter] + pub fn get_top_k(&self) -> Option { + self._as.top_k + } + + #[setter] + pub fn set_top_k(&mut self, k: Option) { + self._as.top_k = k; + } + } + + //////////////////////////////////////////////////////////////////////////////////////////// + // + // OrderBy + // + //////////////////////////////////////////////////////////////////////////////////////////// + + /// One Top-K order key of a :class:`Statement`, set with + /// :meth:`Statement.set_order_by` or :meth:`Statement.then_order_by`. + #[gen_stub_pyclass(module = "_native")] + #[pyclass(frozen, skip_from_py_object, name = "OrderBy", module = "_native")] + #[derive(Clone)] + pub struct OrderBy { + _as: aerospike_core::OrderBy, + } + + #[gen_stub_pymethods] + #[pymethods] + impl OrderBy { + /// Name of the returned bin the key reads. + #[getter] + pub fn bin_name(&self) -> String { + self._as.bin_name.clone() + } + + /// Declared type of the key. + #[getter] + pub fn data_type(&self) -> BinDataType { + self._as.order_type.into() + } + + /// Sort direction. + #[getter] + pub fn order(&self) -> Order { + self._as.direction.into() + } + + fn __eq__(&self, other: &Bound<'_, PyAny>) -> bool { + other.cast::().is_ok_and(|o| o.get()._as == self._as) + } + + fn __hash__(&self) -> u64 { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let key = &self._as; + let mut hasher = DefaultHasher::new(); + (&key.bin_name, key.order_type, key.direction).hash(&mut hasher); + hasher.finish() + } + + fn __repr__(&self, py: Python<'_>) -> PyResult { + Ok(format!( + "OrderBy({}, {}, {})", + self._as.bin_name.as_str().into_pyobject(py)?.repr()?, + Bound::new(py, self.data_type())?.as_any().repr()?, + Bound::new(py, self.order())?.as_any().repr()?, + )) + } } //////////////////////////////////////////////////////////////////////////////////////////// diff --git a/src/lib.rs b/src/lib.rs index 65dee3a3..e75ad24a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5212,13 +5212,19 @@ fn _native(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/src/record.rs b/src/record.rs index 61a5cbeb..4c665797 100644 --- a/src/record.rs +++ b/src/record.rs @@ -15,19 +15,26 @@ use std::collections::hash_map::DefaultHasher; use std::collections::HashMap; +use std::ffi::{c_int, c_void, CStr}; use std::fmt; use std::hash::{Hash, Hasher}; use pyo3::basic::CompareOp; -use pyo3::exceptions::{PyIndexError, PyValueError}; +use pyo3::buffer::{ElementType, PyUntypedBuffer}; +use pyo3::exceptions::{PyBufferError, PyIndexError, PyValueError}; use pyo3::exceptions::PyTypeError; -use pyo3::types::{PyBool, PyByteArray, PyBytes, PyDict, PyList, PySuper, PyTuple}; +use pyo3::ffi; +use pyo3::types::{ + PyBool, PyByteArray, PyBytes, PyDict, PyIterator, PyList, PyMemoryView, PySlice, PySliceMethods, PySuper, + PyTuple, +}; use pyo3::PyTypeInfo; use pyo3::{prelude::*, Borrowed, IntoPyObjectExt}; use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods}; use pyo3_stub_gen::{PyStubType, TypeInfo}; +use aerospike_core::VectorData; /********************************************************************************** @@ -1071,6 +1078,513 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; } } + //////////////////////////////////////////////////////////////////////////////////////////// + // + // Vector + // + //////////////////////////////////////////////////////////////////////////////////////////// + + /// Element type of a :class:`Vector`: how each element is stored on the wire. + #[gen_stub_pyclass_enum(module = "_native")] + #[pyclass(from_py_object, module = "_native")] + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] + pub enum VectorElementType { + /// IEEE 754 half precision. + #[pyo3(name = "FLOAT16")] + Float16, + /// 32-bit signed integer. + #[pyo3(name = "INT32")] + Int32, + /// 32-bit IEEE 754 float. + #[pyo3(name = "FLOAT32")] + Float32, + /// 64-bit IEEE 754 double. + #[pyo3(name = "FLOAT64")] + Float64, + } + + #[pymethods] + impl VectorElementType { + fn __richcmp__(&self, other: &VectorElementType, op: CompareOp) -> PyResult { + match op { + CompareOp::Eq => Ok(self == other), + CompareOp::Ne => Ok(self != other), + _ => Err(pyo3::exceptions::PyNotImplementedError::new_err("Only == and != comparisons are supported")), + } + } + + fn __hash__(&self) -> u64 { + let mut hasher = DefaultHasher::new(); + self.hash(&mut hasher); + hasher.finish() + } + } + + impl VectorElementType { + const fn name(self) -> &'static str { + match self { + VectorElementType::Float16 => "FLOAT16", + VectorElementType::Int32 => "INT32", + VectorElementType::Float32 => "FLOAT32", + VectorElementType::Float64 => "FLOAT64", + } + } + + const fn item_size(self) -> usize { + match self { + VectorElementType::Float16 => 2, + VectorElementType::Int32 | VectorElementType::Float32 => 4, + VectorElementType::Float64 => 8, + } + } + } + + impl From for VectorElementType { + fn from(input: aerospike_core::VectorElementType) -> Self { + match input { + aerospike_core::VectorElementType::Float16 => VectorElementType::Float16, + aerospike_core::VectorElementType::Int32 => VectorElementType::Int32, + aerospike_core::VectorElementType::Float32 => VectorElementType::Float32, + aerospike_core::VectorElementType::Float64 => VectorElementType::Float64, + } + } + } + + /// Compares elements the way Python compares numbers, so `-0.0 == 0.0`. + /// Vectors hold no NaN, which keeps this reflexive. + fn vectors_equal(a: &aerospike_core::Vector, b: &aerospike_core::Vector) -> bool { + match (a.data(), b.data()) { + (VectorData::Float16(x), VectorData::Float16(y)) => { + x.len() == y.len() && x.iter().zip(y).all(|(&p, &q)| p == q || (p | q) & F16_MAGNITUDE == 0) + } + (VectorData::Int32(x), VectorData::Int32(y)) => x == y, + (VectorData::Float32(x), VectorData::Float32(y)) => x == y, + (VectorData::Float64(x), VectorData::Float64(y)) => x == y, + _ => false, + } + } + + const F16_MAGNITUDE: u16 = 0x7fff; + + /// Hashes a vector consistently with [`vectors_equal`]: both zeros hash as `+0.0`. + fn hash_vector(v: &aerospike_core::Vector, state: &mut H) { + v.element_type().code().hash(state); + match v.data() { + VectorData::Float16(d) => d + .iter() + .for_each(|&b| (if b & F16_MAGNITUDE == 0 { 0 } else { b }).hash(state)), + VectorData::Int32(d) => d.hash(state), + VectorData::Float32(d) => d + .iter() + .for_each(|&x| (if x == 0.0 { 0 } else { x.to_bits() }).hash(state)), + VectorData::Float64(d) => d + .iter() + .for_each(|&x| (if x == 0.0 { 0 } else { x.to_bits() }).hash(state)), + } + } + + /// A rejected argument is Python's `ValueError`, as for any constructor: + /// the vector never left the process, so there is no result code to carry. + fn vector_argument_error(e: aerospike_core::Error) -> PyErr { + PyValueError::new_err(e.message().unwrap_or_default().to_owned()) + } + + /// Rounds each element to the nearest `float16`, refusing a finite value + /// that has no finite `float16` form rather than storing an infinity. + fn float16_bits(data: Vec) -> PyResult> { + data.into_iter() + .enumerate() + .map(|(i, x)| { + let h = half::f16::from_f64(x); + if x.is_finite() && !h.is_finite() { + return Err(PyValueError::new_err(format!( + "vector element {i} is out of float16 range: {x}" + ))); + } + Ok(h.to_bits()) + }) + .collect() + } + + fn vector_from_sequence( + data: &Bound<'_, PyAny>, + element_type: VectorElementType, + ) -> PyResult { + let core = match element_type { + VectorElementType::Float16 => aerospike_core::Vector::try_float16(float16_bits(data.extract()?)?), + VectorElementType::Int32 => aerospike_core::Vector::try_int32(data.extract()?), + VectorElementType::Float32 => aerospike_core::Vector::try_float32(data.extract()?), + VectorElementType::Float64 => aerospike_core::Vector::try_float64(data.extract()?), + }; + core.map_err(vector_argument_error) + } + + fn has_buffer(data: &Bound<'_, PyAny>) -> bool { + // SAFETY: `data` is a live object; the call only inspects its type slots. + unsafe { ffi::PyObject_CheckBuffer(data.as_ptr()) == 1 } + } + + const fn is_native_byte_order(order: u8) -> bool { + matches!(order, b'@' | b'=') + || if cfg!(target_endian = "little") { + order == b'<' + } else { + matches!(order, b'>' | b'!') + } + } + + /// The element type a buffer format describes. A byte-swapped format is + /// refused: reading it in place would decode every element wrong. + fn buffer_element_type(format: &CStr) -> Option { + let native = match format.to_bytes() { + [_] => true, + [order, _] => is_native_byte_order(*order), + _ => false, + }; + if !native { + return None; + } + match ElementType::from_format(format) { + ElementType::Float { bytes: 2 } => Some(VectorElementType::Float16), + ElementType::SignedInteger { bytes: 4 } => Some(VectorElementType::Int32), + ElementType::Float { bytes: 4 } => Some(VectorElementType::Float32), + ElementType::Float { bytes: 8 } => Some(VectorElementType::Float64), + _ => None, + } + } + + /// Copies a 1-D buffer's elements, following its stride. + fn buffer_elements(buf: &PyUntypedBuffer) -> Vec { + let count = buf.item_count(); + let base = buf.buf_ptr().cast::().cast_const(); + let stride = buf.strides()[0]; + let size = std::mem::size_of::(); + // SAFETY: while `buf` holds the buffer, its exporter guarantees `count` + // items of `size` bytes at `base + i * stride`; the caller checked the + // item size, and unaligned reads place no alignment demand on `base`. + unsafe { + if stride == size as isize { + let mut out = Vec::::with_capacity(count); + std::ptr::copy_nonoverlapping(base, out.as_mut_ptr().cast::(), count * size); + out.set_len(count); + out + } else { + (0..count) + .map(|i| base.offset(i as isize * stride).cast::().read_unaligned()) + .collect() + } + } + } + + /// Builds a vector from a 1-D buffer such as a numpy array, ``array.array`` + /// or ``memoryview``. The buffer's format selects the element type. + fn vector_from_buffer( + data: &Bound<'_, PyAny>, + element_type: Option, + ) -> PyResult { + // PyO3 refuses an exporter that leaves `strides` null, which the buffer + // protocol allows for contiguous data (ctypes arrays do it); a memoryview + // re-exports the same memory with strides filled in. + let buf = PyUntypedBuffer::get(data) + .or_else(|_| PyUntypedBuffer::get(PyMemoryView::from(data)?.as_any()))?; + if buf.dimensions() != 1 { + return Err(PyValueError::new_err(format!( + "Vector requires a 1-D buffer, got {}-D", + buf.dimensions() + ))); + } + if buf.suboffsets().is_some_and(|s| s[0] >= 0) { + return Err(PyTypeError::new_err("Vector does not read indirect (suboffset) buffers")); + } + let format = buf.format(); + let detected = buffer_element_type(format) + .filter(|t| t.item_size() == buf.item_size()) + .ok_or_else(|| { + PyTypeError::new_err(format!( + "unsupported buffer format '{}' for Vector; use float16, int32, float32 or \ + float64 in native byte order", + format.to_string_lossy() + )) + })?; + if let Some(requested) = element_type { + if requested != detected { + return Err(PyTypeError::new_err(format!( + "element_type {} does not match buffer format '{}'", + requested.name(), + format.to_string_lossy() + ))); + } + } + + let core = match detected { + VectorElementType::Float16 => aerospike_core::Vector::try_float16(buffer_elements(&buf)), + VectorElementType::Int32 => aerospike_core::Vector::try_int32(buffer_elements(&buf)), + VectorElementType::Float32 => aerospike_core::Vector::try_float32(buffer_elements(&buf)), + VectorElementType::Float64 => aerospike_core::Vector::try_float64(buffer_elements(&buf)), + }; + core.map_err(vector_argument_error) + } + + /// A dense numeric vector, stored in a bin as a ``VECTOR`` value and compared + /// with :meth:`FilterExpression.vector_distance`. + /// + /// Build one from a sequence of numbers, with an optional element type + /// (default :attr:`VectorElementType.FLOAT32`), or from a 1-D buffer such as a + /// numpy array, ``array.array`` or ``memoryview``, whose format selects the + /// element type:: + /// + /// embedding = Vector([0.12, 0.98, -0.34]) + /// counts = Vector([3, 0, 7], VectorElementType.INT32) + /// compact = Vector(np.array([0.5, 1.5], dtype=np.float16)) + /// + /// A vector has at least one element, and every floating-point element is + /// finite: an empty sequence, a NaN or an infinity raises :class:`ValueError`. + /// ``FLOAT16`` elements given as Python floats are rounded to the nearest + /// half-precision value. + /// + /// A vector is an immutable, hashable sequence: ``len(v)``, ``v[i]``, slicing + /// and iteration work as they do for a tuple, and two vectors are equal when + /// their element types and elements are equal. It also exposes its elements + /// as a read-only buffer, so ``numpy.asarray(v)`` and ``memoryview(v)`` read + /// them in place without copying. + /// + /// Writing a vector, or sending a vector expression, requires every cluster + /// node to run Aerospike Server 8.2.1 or later; otherwise the client refuses + /// the command before sending it. The server stores a vector only as a whole + /// bin value, not nested in a list or map bin. + #[gen_stub_pyclass(module = "_native")] + #[pyclass(from_py_object, frozen, sequence, module = "_native")] + #[cfg_attr(not(Py_GIL_DISABLED), pyo3(freelist = 1))] + #[derive(Debug, Clone)] + pub struct Vector { + pub(crate) v: aerospike_core::Vector, + } + + impl Vector { + fn element<'py>(&self, py: Python<'py>, index: usize) -> PyResult> { + match self.v.data() { + VectorData::Float16(d) => half::f16::from_bits(d[index]).to_f64().into_bound_py_any(py), + VectorData::Int32(d) => d[index].into_bound_py_any(py), + VectorData::Float32(d) => d[index].into_bound_py_any(py), + VectorData::Float64(d) => d[index].into_bound_py_any(py), + } + } + } + + #[gen_stub_pymethods] + #[pymethods] + impl Vector { + #[new] + #[pyo3(signature = (data, element_type = None))] + pub fn new( + #[gen_stub(override_type( + type_repr = "_native.Vector | typing.Sequence[builtins.float] | collections.abc.Buffer", + imports = ("typing", "builtins", "collections.abc") + ))] + data: &Bound<'_, PyAny>, + element_type: Option, + ) -> PyResult { + if let Ok(existing) = data.cast::() { + let existing = existing.get(); + if element_type.is_some_and(|t| t != existing.v.element_type().into()) { + return Err(PyTypeError::new_err( + "element_type does not match the Vector's element type", + )); + } + return Ok(existing.clone()); + } + // Lists and tuples are the common input and never export a buffer. + if !data.is_instance_of::() && !data.is_instance_of::() && has_buffer(data) { + return Ok(Vector { v: vector_from_buffer(data, element_type)? }); + } + let v = vector_from_sequence(data, element_type.unwrap_or(VectorElementType::Float32))?; + Ok(Vector { v }) + } + + /// Decodes a complete vector wire value: a version byte, an element-type + /// byte, the dimension count as a little-endian ``u32``, then the + /// little-endian elements. Use it for a vector carried as ``bytes``, such + /// as one an application stored in a blob bin. + /// + /// Example:: + /// + /// record = await client.get(key, ["embedding_blob"]) + /// embedding = Vector.from_wire_bytes(record.bins["embedding_blob"]) + /// + /// Args: + /// data: The encoded vector. + /// + /// Returns: + /// The decoded vector. + /// + /// Raises: + /// ValueError: The bytes are not a valid vector wire value. + #[staticmethod] + pub fn from_wire_bytes( + #[gen_stub(override_type(type_repr = "builtins.bytes", imports = ("builtins")))] data: Vec, + ) -> PyResult { + let v = aerospike_core::Vector::from_wire_bytes(&data).map_err(vector_argument_error)?; + Ok(Vector { v }) + } + + /// The element type. + #[getter] + pub fn element_type(&self) -> VectorElementType { + self.v.element_type().into() + } + + /// The number of elements. + #[getter] + pub fn dimensions(&self) -> usize { + self.v.dimensions() + } + + /// The elements as a new list of ``int`` (``INT32``) or ``float`` values. + #[gen_stub(override_return_type( + type_repr = "builtins.list[builtins.int] | builtins.list[builtins.float]", + imports = ("builtins") + ))] + #[getter] + pub fn value<'py>(&self, py: Python<'py>) -> PyResult> { + match self.v.data() { + VectorData::Float16(d) => d + .iter() + .map(|&bits| half::f16::from_bits(bits).to_f64()) + .collect::>() + .into_bound_py_any(py), + VectorData::Int32(d) => d.into_bound_py_any(py), + VectorData::Float32(d) => d.into_bound_py_any(py), + VectorData::Float64(d) => d.into_bound_py_any(py), + } + } + + fn __len__(&self) -> usize { + self.v.dimensions() + } + + #[gen_stub(override_return_type( + type_repr = "builtins.int | builtins.float | builtins.list[builtins.int] | builtins.list[builtins.float]", + imports = ("builtins") + ))] + fn __getitem__<'py>(&self, py: Python<'py>, index: &Bound<'py, PyAny>) -> PyResult> { + let len = self.v.dimensions() as isize; + if let Ok(slice) = index.cast::() { + let bounds = slice.indices(len)?; + let elements = (0..bounds.slicelength as isize) + .map(|i| self.element(py, (bounds.start + i * bounds.step) as usize)) + .collect::>>()?; + return PyList::new(py, elements).map(Bound::into_any); + } + let i: isize = index.extract()?; + let position = if i < 0 { i + len } else { i }; + if !(0..len).contains(&position) { + return Err(PyIndexError::new_err("Vector index out of range")); + } + self.element(py, position as usize) + } + + #[gen_stub(override_return_type( + type_repr = "typing.Iterator[builtins.int | builtins.float]", + imports = ("typing", "builtins") + ))] + fn __iter__<'py>(&self, py: Python<'py>) -> PyResult> { + self.value(py)?.try_iter() + } + + /// Exposes the elements as a read-only buffer in native byte order, so + /// ``memoryview(v)`` and ``numpy.asarray(v)`` read them without copying. + #[gen_stub(skip)] + unsafe fn __getbuffer__(slf: Bound<'_, Self>, view: *mut ffi::Py_buffer, flags: c_int) -> PyResult<()> { + if view.is_null() { + return Err(PyBufferError::new_err("buffer view is null")); + } + if flags & ffi::PyBUF_WRITABLE == ffi::PyBUF_WRITABLE { + return Err(PyBufferError::new_err("Vector is read-only")); + } + let v = &slf.get().v; + let (data, item_size, format): (*const c_void, usize, &'static CStr) = match v.data() { + VectorData::Float16(d) => (d.as_ptr().cast(), 2, c"e"), + VectorData::Int32(d) => (d.as_ptr().cast(), 4, c"i"), + VectorData::Float32(d) => (d.as_ptr().cast(), 4, c"f"), + VectorData::Float64(d) => (d.as_ptr().cast(), 8, c"d"), + }; + let count = v.dimensions(); + // SAFETY: `view` is non-null and owned by the caller. The elements stay + // valid and unchanged for the view's lifetime: the Vector is immutable + // and `obj` keeps it alive. The shape box is freed in `__releasebuffer__`. + unsafe { + let shape = Box::into_raw(Box::new(count as ffi::Py_ssize_t)); + (*view).buf = data.cast_mut(); + (*view).len = (count * item_size) as ffi::Py_ssize_t; + (*view).readonly = 1; + (*view).itemsize = item_size as ffi::Py_ssize_t; + (*view).format = if flags & ffi::PyBUF_FORMAT == ffi::PyBUF_FORMAT { + format.as_ptr().cast_mut() + } else { + std::ptr::null_mut() + }; + (*view).ndim = 1; + (*view).shape = if flags & ffi::PyBUF_ND == ffi::PyBUF_ND { shape } else { std::ptr::null_mut() }; + (*view).strides = if flags & ffi::PyBUF_STRIDES == ffi::PyBUF_STRIDES { + &mut (*view).itemsize + } else { + std::ptr::null_mut() + }; + (*view).suboffsets = std::ptr::null_mut(); + (*view).internal = shape.cast(); + (*view).obj = slf.into_any().into_ptr(); + } + Ok(()) + } + + #[gen_stub(skip)] + unsafe fn __releasebuffer__(&self, view: *mut ffi::Py_buffer) { + // SAFETY: `internal` holds the shape box `__getbuffer__` allocated for this view. + drop(unsafe { Box::from_raw((*view).internal.cast::()) }); + } + + fn __richcmp__(&self, other: &Bound<'_, PyAny>, op: CompareOp) -> bool { + let equal = other.cast::().is_ok_and(|other| vectors_equal(&self.v, &other.get().v)); + match op { + CompareOp::Eq => equal, + CompareOp::Ne => !equal, + _ => false, + } + } + + fn __hash__(&self) -> u64 { + let mut hasher = DefaultHasher::new(); + hash_vector(&self.v, &mut hasher); + hasher.finish() + } + + fn __repr__(&self) -> String { + let elements = match self.v.data() { + VectorData::Float16(d) => format!( + "{:?}", + d.iter().map(|&b| half::f16::from_bits(b).to_f64()).collect::>() + ), + VectorData::Int32(d) => format!("{d:?}"), + VectorData::Float32(d) => format!("{d:?}"), + VectorData::Float64(d) => format!("{d:?}"), + }; + let element_type: VectorElementType = self.v.element_type().into(); + format!("Vector({elements}, VectorElementType.{})", element_type.name()) + } + } + + // The buffer slots take raw `Py_buffer` pointers, which the stub generator + // cannot type; these are the Python-level methods CPython derives from them. + pyo3_stub_gen::inventory::submit! { + pyo3_stub_gen::derive::gen_methods_from_python! { + r#" + class Vector: + def __buffer__(self, flags: int) -> memoryview: ... + def __release_buffer__(self, buffer: memoryview) -> None: ... + "# + } + } + //////////////////////////////////////////////////////////////////////////////////////////// // // SpecialValue (CDT range / value boundaries) @@ -1161,6 +1675,9 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; /// HLL value HLL(Vec), + /// Dense numeric vector (see :class:`Vector`). + Vector(aerospike_core::Vector), + /// CDT boundary markers (see :class:`SpecialValue`). CdtSpecial(SpecialValue), } @@ -1179,6 +1696,7 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; PythonValue::String(ref val) | PythonValue::GeoJSON(ref val) => val.hash(state), PythonValue::Blob(ref val) | PythonValue::HLL(ref val) => val.hash(state), PythonValue::List(ref val) => val.hash(state), + PythonValue::Vector(ref val) => hash_vector(val, state), PythonValue::HashMap(_) | PythonValue::OrderedMap(_) | PythonValue::SortedMap(_) => { @@ -1201,6 +1719,7 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; PythonValue::GeoJSON(ref val) => format!("GeoJSON('{}')", val), PythonValue::Blob(ref val) => format!("{:?}", val), PythonValue::HLL(ref val) => format!("HLL('{:?}')", val), + PythonValue::Vector(ref val) => val.to_string(), PythonValue::List(ref val) => format!("{:?}", val), PythonValue::HashMap(ref val) => format!("{:?}", val), PythonValue::SortedMap(ref val) => format!("{:?}", val), @@ -1277,6 +1796,7 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; Ok(geo.into_pyobject(py).map(|v| v.into_any()).unwrap()) } PythonValue::HLL(b) => Ok(HLL::new(b).into_pyobject(py).map(|v| v.into_any()).unwrap()), + PythonValue::Vector(v) => Ok(Vector { v }.into_pyobject(py).map(|v| v.into_any()).unwrap()), PythonValue::CdtSpecial(s) => { Ok(s.into_pyobject(py).map(|v| v.into_any()).unwrap()) } @@ -1293,6 +1813,12 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; return Ok(PythonValue::Nil); } + // Ahead of the sequence extraction below, which would otherwise read a + // Vector as a list of numbers. + if let Ok(vector) = obj.cast::() { + return Ok(PythonValue::Vector(vector.get().v.clone())); + } + if let Ok(sv) = ::extract(obj) { return Ok(PythonValue::CdtSpecial(sv)); } @@ -1440,6 +1966,7 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; } PythonValue::GeoJSON(gj) => aerospike_core::Value::GeoJson(gj), PythonValue::HLL(b) => aerospike_core::Value::Hll(b), + PythonValue::Vector(v) => aerospike_core::Value::Vector(v), PythonValue::CdtSpecial(s) => match s { SpecialValue::Null => aerospike_core::Value::Nil, SpecialValue::Infinity => aerospike_core::Value::Infinity, @@ -1504,6 +2031,7 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; } aerospike_core::Value::GeoJson(gj) => PythonValue::GeoJSON(gj), aerospike_core::Value::Hll(b) => PythonValue::HLL(b), + aerospike_core::Value::Vector(v) => PythonValue::Vector(v), aerospike_core::Value::Infinity => PythonValue::CdtSpecial(SpecialValue::Infinity), aerospike_core::Value::Wildcard => PythonValue::CdtSpecial(SpecialValue::Wildcard), aerospike_core::Value::KeyValueList(kvl) => {