diff --git a/docs/modules/ROOT/pages/how-to/vectors.adoc b/docs/modules/ROOT/pages/how-to/vectors.adoc index 828a51e..b64d9f2 100644 --- a/docs/modules/ROOT/pages/how-to/vectors.adoc +++ b/docs/modules/ROOT/pages/how-to/vectors.adoc @@ -459,6 +459,8 @@ These rules hold for the quantization arguments of the method: * `encoding="INT8"` is for a `BINARY` property that already holds INT8 bytes. Set `quantization` to `"NONE"` with it. The call raises `ValueError` when both are `INT8`, because the vector would be quantized twice. * A bucket that holds fewer vectors than `pq_clusters` is too small to train a codebook. The index still builds: the engine skips PQ for that bucket and searches it exactly, so `find_nearest_approximate` returns the exact neighbours with exact scores from it. A bucket with `pq_clusters` vectors or more is trained and answers approximately. +On the `arcadedb` (native) package, `find_nearest_approximate` hands the engine its query as a float32 array by address and reads the hits from one buffer that Python provides, with no JSON text on either side. The call returns the same list as before: the same records in the same order, the same RID texts, and scores of the same type with the same bits. A query crosses this way when it is a one-dimensional `float16`, `float32` or `float64` NumPy array, or a list or tuple of `float` and `int` values (an `int` beyond 2**53 is not one), and has no NaN or infinity in it; an `allowed_rids` filter crosses this way when every entry is a plain `#bucket:position` text or a record. Anything else keeps the JSON text route, which returns the same hits and only takes longer: a generator, a NumPy scalar or `None` in a list, an integer array, an empty query, a `float64` value that float32 cannot hold. `ARCADEDB_NATIVE_VINDEX_BIN=0` in the environment turns the address route off. A native library that is older than this change falls back to the text route by itself. `find_nearest` and the other searches are statements and are not changed by this: they already send a NumPy query as float32 and read their rows as a result set. + == Rebuild the graph Inserts, updates and deletes collect as pending changes. When enough of them have collected, the engine rebuilds the graph in the background, and searches keep running. How many is enough comes from the database setting `arcadedb.vectorIndex.rebuildGraphRatio` and from `mutations_before_rebuild`; see xref:how-to/databases-and-queries.adoc#database-settings[Database settings] for how to change the ratio. diff --git a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java index b1d3e7f..1fd844e 100644 --- a/native/src/main/java/org/humemai/nativelib/ArcadeNative.java +++ b/native/src/main/java/org/humemai/nativelib/ArcadeNative.java @@ -681,6 +681,11 @@ public static CCharPointer insertColp(final IsolateThread thread, final long h, private static final ThreadLocal LAST_ERROR = new ThreadLocal<>(); + /** Records the error of an entry point that answers with a long (-1), for arcadedb_n_last_error on the same thread. */ + static void recordError(final Throwable t) { + LAST_ERROR.set(errorJson(t)); + } + /** * A command (flags bit 0 selects command, else query) whose first rows are pulled but not encoded (as flags bit 1 of arcadedb_n_exec). * Returns the cursor id (> 0, read with arcadedb_n_next, closed with arcadedb_n_rs_close or arcadedb_n_rs_close_many), 0 when the diff --git a/native/src/main/java/org/humemai/nativelib/NativeExtras.java b/native/src/main/java/org/humemai/nativelib/NativeExtras.java index d230919..7a7826f 100644 --- a/native/src/main/java/org/humemai/nativelib/NativeExtras.java +++ b/native/src/main/java/org/humemai/nativelib/NativeExtras.java @@ -359,6 +359,7 @@ private static String q(final String text) { * ops: metadata {type, prop} | approximate {type, prop, vector, k, allowed} | build {type, prop}. The index is the one the type has on the * property; its bucket indexes are searched in turn (the caller merges). metadata describes the first bucket index, as the JPype wrapper does. * build rebuilds the graph of every bucket index now (buildVectorGraphNow), as the JPype wrapper's build_graph_now() does (#162). + * approximate has a binary twin, arcadedb_n_vindex_bin (below); this one stays for the queries that entry cannot carry. */ @CEntryPoint(name = "arcadedb_n_vindex") @SuppressWarnings("unchecked") @@ -367,15 +368,7 @@ public static CCharPointer vindex(final IsolateThread thread, final long dbh, fi final Database db = ArcadeNative.db(dbh); final Map a = args(argsJson); final String type = (String) a.get("type"), prop = (String) a.get("prop"); - final TypeIndex typeIndex = db.getSchema().getType(type).getPolymorphicIndexByProperties(prop); - if (typeIndex == null) - throw new IllegalArgumentException("No vector index on " + type + "[" + prop + "]"); - final List parts = new ArrayList<>(); - for (final IndexInternal index : typeIndex.getIndexesOnBuckets()) - if (index instanceof LSMVectorIndex lsm) - parts.add(lsm); - if (parts.isEmpty()) - throw new IllegalArgumentException("Underlying index is not an LSM vector index"); + final List parts = vectorParts(db, type, prop); final StringBuilder sb = new StringBuilder("{\"ok\":true"); switch (ArcadeNative.str(op)) { case "metadata" -> { @@ -413,15 +406,11 @@ public static CCharPointer vindex(final IsolateThread thread, final long dbh, fi } sb.append(",\"hits\":["); boolean firstHit = true; - for (final LSMVectorIndex lsm : parts) { - final List> hits = allowed != null ? lsm.findNeighborsFromVectorApproximate(query, k, allowed) - : lsm.findNeighborsFromVectorApproximate(query, k); - for (final Pair hit : hits) { - if (!firstHit) - sb.append(','); - firstHit = false; - sb.append("[\"").append(hit.getFirst()).append("\",").append(hit.getSecond()).append(']'); - } + for (final Pair hit : approximateHits(parts, query, k, allowed)) { + if (!firstHit) + sb.append(','); + firstHit = false; + sb.append("[\"").append(hit.getFirst()).append("\",").append(hit.getSecond()).append(']'); } sb.append(']'); } @@ -437,6 +426,85 @@ public static CCharPointer vindex(final IsolateThread thread, final long dbh, fi } } + /** The bucket indexes of the vector index that `type` has on `prop`. */ + private static List vectorParts(final Database db, final String type, final String prop) { + final TypeIndex typeIndex = db.getSchema().getType(type).getPolymorphicIndexByProperties(prop); + if (typeIndex == null) + throw new IllegalArgumentException("No vector index on " + type + "[" + prop + "]"); + final List parts = new ArrayList<>(); + for (final IndexInternal index : typeIndex.getIndexesOnBuckets()) + if (index instanceof LSMVectorIndex lsm) + parts.add(lsm); + if (parts.isEmpty()) + throw new IllegalArgumentException("Underlying index is not an LSM vector index"); + return parts; + } + + /** The PQ-scored hits of every bucket index in turn, one list (up to k hits per bucket index; the caller merges). */ + private static List> approximateHits(final List parts, final float[] query, final int k, + final Set allowed) { + final List> all = new ArrayList<>(); + for (final LSMVectorIndex lsm : parts) + all.addAll(allowed != null ? lsm.findNeighborsFromVectorApproximate(query, k, allowed) + : lsm.findNeighborsFromVectorApproximate(query, k)); + return all; + } + + /** + * The "approximate" op of arcadedb_n_vindex with no JSON on either side: no number text for the query, no hit text for the answer. + * type and prop are C strings. The query is `dims` float32 values at `query` (the engine narrows the JSON route's numbers to float32 + * as well, so the search sees the same floats). `allowed` holds `allowedPairs` (bucket, position) pairs of int64, or is NULL with 0 + * for no filter. `out` has room for `capHits` hits of three int64 words each: the RID's bucket, its position, and the score as the bits of + * a double. The score is the double that the JSON route's number text of the same float parses to (Float.toString, then + * Double.parseDouble: the shortest decimal of the float), so a caller reads the same bits from either route. + * + * Returns the number of hits n >= 0. When n <= capHits the n hits are in `out`; when n is larger nothing is written, and the caller calls + * again with room for n. -1 is an error, its text is read with arcadedb_n_last_error on the same thread (no reply text: no malloc, no + * arcadedb_free and no JSON decode in the caller). + */ + @CEntryPoint(name = "arcadedb_n_vindex_bin") + public static long vindexBin(final IsolateThread thread, final long dbh, final CCharPointer type, final CCharPointer prop, final int k, + final CFloatPointer query, final long dims, final CCharPointer allowed, final long allowedPairs, final CCharPointer out, + final long capHits) { + try { + if (dims < 0 || dims > Integer.MAX_VALUE || allowedPairs < 0 || capHits < 0 || capHits > Integer.MAX_VALUE / 24) + throw new IllegalArgumentException( + "vector search arguments out of range (dims " + dims + ", allowed " + allowedPairs + ", capacity " + capHits + ")"); + if (dims > 0 && query.isNull()) + throw new IllegalArgumentException("no query vector"); + final Database db = ArcadeNative.db(dbh); + final List parts = vectorParts(db, ArcadeNative.str(type), ArcadeNative.str(prop)); + final float[] q = new float[(int) dims]; + for (int i = 0; i < q.length; i++) + q[i] = query.read(i); + Set allowedSet = null; + if (allowedPairs > 0) { + if (allowed.isNull()) + throw new IllegalArgumentException("no allowed RIDs"); + final Pointer pairs = (Pointer) allowed; + allowedSet = new HashSet<>(); + for (long i = 0; i < allowedPairs; i++) + allowedSet.add( + new RID((int) pairs.readLong(WordFactory.signed(16L * i)), pairs.readLong(WordFactory.signed(16L * i + 8)))); + } + final List> hits = approximateHits(parts, q, k, allowedSet); + final int n = hits.size(); + if (n > capHits) + return n; + final Pointer o = (Pointer) out; + for (int i = 0; i < n; i++) { + final Pair hit = hits.get(i); + o.writeLong(24 * i, hit.getFirst().getBucketId()); + o.writeLong(24 * i + 8, hit.getFirst().getPosition()); + o.writeLong(24 * i + 16, Double.doubleToRawLongBits(Double.parseDouble(Float.toString(hit.getSecond())))); + } + return n; + } catch (final Throwable t) { + ArcadeNative.recordError(t); + return -1; + } + } + // ---- GraphBatch private record BatchEntry(GraphBatch batch, Database db) { diff --git a/src/arcadedb/_native/lib.py b/src/arcadedb/_native/lib.py index 9f7474d..b9ab272 100644 --- a/src/arcadedb/_native/lib.py +++ b/src/arcadedb/_native/lib.py @@ -37,6 +37,9 @@ # ARCADEDB_NATIVE_GB_BIN: the float32 vectors of a GraphBatch op (create_vertices rows, new_edges props) go through arcadedb_n_gb_bin in a float32 side buffer, not as number text GB_BIN = os.environ.get("ARCADEDB_NATIVE_GB_BIN", "1") != "0" +# ARCADEDB_NATIVE_VINDEX_BIN: the query vector of find_nearest_approximate crosses as a float32 array by address and the hits come back in one packed buffer (arcadedb_n_vindex_bin), not as a JSON list of numbers and a JSON list of [rid, score] +VINDEX_BIN = os.environ.get("ARCADEDB_NATIVE_VINDEX_BIN", "1") != "0" + # the directory the engine log, logging.properties and the crash files go to (None until the library is loaded) LOG_DIR_USED = None @@ -255,6 +258,7 @@ class _Params(ctypes.Structure): "arcadedb_n_async": [P, L, S, S], "arcadedb_n_async_bin": [P, L, S, S, P, N], "arcadedb_n_vindex": [P, L, S, S], + "arcadedb_n_vindex_bin": [P, L, S, S, N, P, L, P, L, P, L], "arcadedb_n_import": [P, L, S], "arcadedb_n_export": [P, L, S, S, N, S], "arcadedb_n_cb": [P, S, S], @@ -262,9 +266,10 @@ class _Params(ctypes.Structure): "arcadedb_n_set_upcall": [P, P], "arcadedb_free": [P, P], } - # the entry points that answer with a number, not with the address of a reply text: a cursor id (arcadedb_n_cmd) and the - # transaction state (arcadedb_n_tx); -1 is an error for both, its text is read with arcadedb_n_last_error - returns_long = ("arcadedb_n_cmd", "arcadedb_n_tx") + # the entry points that answer with a number, not with the address of a reply text: a cursor id (arcadedb_n_cmd), the + # transaction state (arcadedb_n_tx) and the hit count of arcadedb_n_vindex_bin; -1 is an error for all three, its text is + # read with arcadedb_n_last_error + returns_long = ("arcadedb_n_cmd", "arcadedb_n_tx", "arcadedb_n_vindex_bin") for name, argtypes in sigs.items(): fn = getattr(lib, name, None) if fn is None: # a library built before this entry point existed @@ -280,7 +285,7 @@ class _Params(ctypes.Structure): ) else ctypes.c_longlong if name in returns_long else ctypes.c_void_p ) - global BINARY_PARAMS, CMD_PATH, CLOSE_BATCH, INSERT_ROWS, EXEC_MANY, NEXT_BIN, GB_BIN, TS_BIN + global BINARY_PARAMS, CMD_PATH, CLOSE_BATCH, INSERT_ROWS, EXEC_MANY, NEXT_BIN, GB_BIN, TS_BIN, VINDEX_BIN if not hasattr( lib, "arcadedb_n_exec_bin" ): # a library built before the binary path @@ -305,6 +310,10 @@ class _Params(ctypes.Structure): lib, "arcadedb_n_async_bin" ): # a library built before the address table of append_samples TS_BIN = False + if not hasattr( + lib, "arcadedb_n_vindex_bin" + ): # a library built before the binary approximate search + VINDEX_BIN = False _isolate = isolate _creator = threading.get_ident() _local.thread = _Thread(thread, owned=False) @@ -505,8 +514,27 @@ def tx(handle, op) -> int: return lib.arcadedb_n_tx(th, handle, op) +def vindex_bin( + handle, type_name, prop, k, query, dims, allowed, pairs, out, cap +) -> int: + """``arcadedb_n_vindex_bin``: the approximate search of a vector index with no JSON on either side. + + ``type_name`` and ``prop`` are bytes; ``query`` is the address of ``dims`` float32 values; ``allowed`` is the address of ``pairs`` + (bucket, position) int64 pairs, or 0 with 0 pairs for no filter; ``out`` is the address of room for ``cap`` hits of three int64 + words (bucket, position, the score as the bits of a double). Returns the hit count n: the n hits are in ``out`` when ``n <= cap``, + and nothing is written when it is larger (call again with room for n). -1 is an error, read with :func:`last_error`. + The caller keeps every buffer alive until the call returns. + """ + lib = _lib or _load() + t = _local.__dict__.get("thread") + th = t.handle if t is not None else _thread_handle() + return lib.arcadedb_n_vindex_bin( + th, handle, type_name, prop, k, query, dims, allowed, pairs, out, cap + ) + + def last_error() -> str: - """The error of the last failed :func:`cmd` or :func:`tx` on this thread, as the JSON text of an error reply.""" + """The error of the last failed :func:`cmd`, :func:`tx` or :func:`vindex_bin` on this thread, as the JSON text of an error reply.""" return call("arcadedb_n_last_error") diff --git a/src/arcadedb/_native/schema_objects.py b/src/arcadedb/_native/schema_objects.py index 5f69e0d..af14905 100644 --- a/src/arcadedb/_native/schema_objects.py +++ b/src/arcadedb/_native/schema_objects.py @@ -13,6 +13,24 @@ _PROPERTY = re.compile(r"name: (.+?), type: ([A-Z_]+)") +# find_nearest_approximate crosses in arcadedb_n_vindex_bin. The answer buffer is sized for this many bucket indexes of the type at +# first (each answers with up to k hits); a type with more of them costs one more crossing, and the larger size is remembered. +_VINDEX_PARTS = 8 +# an answer buffer larger than this is not asked for: the JSON route answers +_VINDEX_MAX_BYTES = 1 << 30 +# int64 words per hit: bucket, position, the score as the bits of a double +_HIT_WORDS = 3 +_RID_PAIRS = re.compile(r"#\d{1,9}:\d{1,18}(?:;#\d{1,9}:\d{1,18})*") +_DIGITS = re.compile(r"\d+") +_NUMBERS = frozenset((float, int)) +_FLOAT_DTYPES = ( + (_np.dtype(_np.float16), _np.dtype(_np.float32), _np.dtype(_np.float64)) + if _np is not None + else () +) +# an int up to this is a double, so float32 rounds it as Java's floatValue() of the parsed JSON number does +_EXACT_INT = 1 << 53 + def _name_type(entry): """(name, type) of an entry of a ``schema:types`` row's ``properties``: a dict since #195 (the codec writes a nested row @@ -98,6 +116,7 @@ class _IndexShim: def __init__(self, database, type_name, prop, name): self._db, self._type, self._prop, self._name = database, type_name, prop, name + self._vindex_parts = _VINDEX_PARTS def getName(self): return self._name @@ -170,6 +189,59 @@ def _rid_text(value) -> str: return get_rid() if get_rid is not None else str(value) +def _vindex_query(query_vector): + """The query of an approximate search as a contiguous float32 array, for ``arcadedb_n_vindex_bin``; None when the JSON route must + carry it. + + The engine searches with float32 whatever it is given: the JSON route sends doubles and Java narrows them with ``floatValue()``, which + rounds to nearest even, as numpy's cast to float32 does, so the library reads the same floats from the address. What goes by + address: a one-dimensional ``float16``, ``float32`` or ``float64`` array of the native byte order, and a list or tuple of ``float`` + and ``int`` (an int beyond 2**53 is not a double and stays JSON). Everything else keeps the JSON route (it reads whatever the + object's ``tolist()`` or ``list()`` gives and reports what the engine makes of it): a generator, any other array or sequence, an + entry that is not a number, a bool, a NaN or an infinity (also one that float32 overflows to), and an empty query. + """ + if _np is None: + return None + if type(query_vector) is _np.ndarray: + if query_vector.ndim != 1 or query_vector.dtype not in _FLOAT_DTYPES: + return None + source = query_vector + elif type(query_vector) in (list, tuple): + kinds = set(map(type, query_vector)) + if not kinds or not kinds <= _NUMBERS: + return None + if int in kinds and any( + abs(v) > _EXACT_INT for v in query_vector if type(v) is int + ): + return None + source = query_vector + else: + return None + try: + with _np.errstate(over="ignore", invalid="ignore"): + f32 = _np.asarray(source, dtype=_np.float32) + except (OverflowError, ValueError, TypeError): + return None + if f32.size == 0 or not _np.isfinite(f32).all(): + return None + return _np.require(f32, requirements=["C", "A"]) + + +def _vindex_allowed(texts): + """The RIDs of an approximate search's filter (``"#bucket:position"`` texts) as an int64 array of (bucket, position) pairs, or None + when one of them is not that plain form (the JSON route hands the text to the engine's own RID parser). + """ + try: + joined = ";".join(texts) + except TypeError: + return None + numbers = _DIGITS.findall(joined) if _RID_PAIRS.fullmatch(joined) else () + # one text that holds a ';' is two RIDs to the pattern but one bad RID to the engine: the count tells them apart + if not numbers or len(numbers) != 2 * len(texts): + return None + return _np.array(numbers, dtype=_np.int64) + + def _query_vector_param(query_vector): """The query vector as the one parameter the library takes it in. @@ -331,6 +403,71 @@ def _run_search_without_vector(self, query_vector, k, allowed_rids, ef_search): def _run_approximate(self, query_vector, k, allowed_rids): """The PQ-scored search of the engine (no statement runs it): the engine returns (RID, distance), the records are loaded here.""" from . import lib + + vector = _vindex_query(query_vector) if lib.VINDEX_BIN else None + k = int(k) + allowed = [_rid_text(r) for r in allowed_rids] if allowed_rids else None + hits = None if vector is None else self._approximate_packed(vector, k, allowed) + if hits is None: + hits = self._approximate_json(query_vector, k, allowed) + out = [(self._database.lookup_by_rid(rid), score) for rid, score in hits] + return self._sort_results(out)[:k] + + def _approximate_packed(self, vector, k, allowed): + """``[(rid text, score)]`` from ``arcadedb_n_vindex_bin``, or None when the JSON route has to carry the call (a k that is not a C + int, a filter that is not plain ``#bucket:position`` texts, an answer buffer of more than a gibibyte, or one that the engine keeps + outgrowing). The library fills one buffer of int64 words, three per hit, and answers with their number; the query, the filter and + the buffer are numpy arrays that live until the call returns. + """ + from . import lib + from .database import _error + + if not -(1 << 31) <= k < (1 << 31): + return None + pairs, count = None, 0 + if allowed: + pairs = _vindex_allowed(allowed) + if pairs is None: + return None + count = len(pairs) // 2 + shim = self._shim + room = max(k, 1) + for _ in range(3): + cap = room * shim._vindex_parts + if cap * _HIT_WORDS * 8 > _VINDEX_MAX_BYTES: + return None + out = _np.empty(cap * _HIT_WORDS, dtype=_np.int64) + n = lib.vindex_bin( + self._database._h, + lib.bc(shim._type), + lib.bc(shim._prop), + k, + vector.ctypes.data, + vector.size, + pairs.ctypes.data if count else 0, + count, + out.ctypes.data, + cap, + ) + if n < 0: + raise _error("Approximate search failed", json.loads(lib.last_error())) + if n <= cap: + words, end = out[: n * _HIT_WORDS], n * _HIT_WORDS + rids = [ + f"#{b}:{p}" + for b, p in zip( + words[0::_HIT_WORDS].tolist(), words[1::_HIT_WORDS].tolist() + ) + ] + scores = out.view(_np.float64)[2:end:_HIT_WORDS].tolist() + return list(zip(rids, scores)) + # more hits than the buffer holds: nothing was written; ask again with room for them, and keep that size + shim._vindex_parts = -(-n // room) + return None + + def _approximate_json(self, query_vector, k, allowed): + """``[(rid text, score)]`` through the JSON text of ``arcadedb_n_vindex``: the route for every call the binary entry cannot carry.""" + from . import lib from .extras import _call, _enc vector = ( @@ -349,21 +486,13 @@ def _run_approximate(self, query_vector, k, allowed_rids): "type": self._shim._type, "prop": self._shim._prop, "vector": vector, - "k": int(k), - "allowed": ( - [_rid_text(r) for r in allowed_rids] - if allowed_rids - else None - ), + "k": k, + "allowed": allowed, } ) ), ) - out = [ - (self._database.lookup_by_rid(rid), float(distance)) - for rid, distance in reply["hits"] - ] - return self._sort_results(out)[:k] + return [(rid, float(distance)) for rid, distance in reply["hits"]] def find_nearest_approximate(self, query_vector, k=10, allowed_rids=None): """PQ approximate search over PRODUCT quantization, the engine's own (``findNeighborsFromVectorApproximate``).""" diff --git a/tests/conftest.py b/tests/conftest.py index 2920c28..e375c95 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -113,6 +113,8 @@ collect_ignore.append("test_native_vertex_vector_crossings.py") # the native columnar fast path (arcadedb._native.bulk.insert_columns_fast) with a rows-dropped fault in it, found by the ingest-integrity check (#336) collect_ignore.append("test_native_ingest_integrity_faults.py") + # find_nearest_approximate through arcadedb_n_vindex_bin against the JSON route of arcadedb_n_vindex, observed on the native library + collect_ignore.append("test_native_vindex_bin.py") if not os.path.isfile( os.path.join(os.path.dirname(os.path.dirname(__file__)), "docs", "index.md") ): diff --git a/tests/test_native_vindex_bin.py b/tests/test_native_vindex_bin.py new file mode 100644 index 0000000..050f833 --- /dev/null +++ b/tests/test_native_vindex_bin.py @@ -0,0 +1,301 @@ +"""Native wheel: ``find_nearest_approximate`` through ``arcadedb_n_vindex_bin`` against the JSON route it replaces. + +The query vector used to cross as a JSON list of d numbers and the hits came back as a JSON list of ``[rid, score]``. Now a float32 query +crosses by address and the hits come back in one buffer. The answer must be the same list: the same records in the same order, the same +RID texts, and scores of the same type with the same bits (a score is the double that the number text of the engine's float parses to). +The JSON route stays for what the entry cannot carry, and ``lib.VINDEX_BIN`` (``ARCADEDB_NATIVE_VINDEX_BIN=0``) selects it, which is how +these tests run the old route beside the new one in one process. + +Whether the new route is faster is a timed A/B, not these tests. What the Python half does with its arguments and its answer, with a +stand-in for the library, is in ``test_native_vindex_glue.py``. This file needs the native library, so conftest.py does not collect it on +the JVM backend. +""" + +import struct + +import numpy as np +import pytest + +import arcadedb +from arcadedb._native import lib +from arcadedb.exceptions import ArcadeDBError + +DIM, N, BUCKETS = ( + 16, + 1200, + 4, +) # 300 vectors a bucket: above the 64 PQ clusters, so the search is approximate + + +@pytest.fixture(scope="module") +def pq_index(tmp_path_factory): + path = str(tmp_path_factory.mktemp("vindex_bin") / "db") + db = arcadedb.create_database(path) + db.command("sql", f"CREATE VERTEX TYPE Doc BUCKETS {BUCKETS}") + db.command("sql", "CREATE PROPERTY Doc.id INTEGER") + db.command("sql", "CREATE PROPERTY Doc.embedding ARRAY_OF_FLOATS") + index = db.create_vector_index( + "Doc", + "embedding", + dimensions=DIM, + id_property="id", + quantization="PRODUCT", + pq_subspaces=4, + pq_clusters=64, + build_graph_now=False, + ) + vectors = np.random.default_rng(11).standard_normal((N, DIM)).astype(np.float32) + with db.transaction(): + for i in range(N): + db.command( + "sql", "INSERT INTO Doc SET id = ?, embedding = ?", i, vectors[i] + ) + index.build_graph_now() + index.vectors = vectors + index.database = db + yield index + db.drop() + + +class _Calls: + """The approximate searches that reach the library, by entry point.""" + + def __init__(self, monkeypatch): + self.binary, self.text = [], [] + real_bin, real_call = lib.vindex_bin, lib.call + + def vindex_bin(*args): + n = real_bin(*args) + self.binary.append( + { + "k": args[3], + "dims": args[5], + "pairs": args[7], + "cap": args[9], + "n": n, + } + ) + return n + + def call(name, *args): + if name == "arcadedb_n_vindex" and args[1] == b"approximate": + self.text.append(args) + return real_call(name, *args) + + monkeypatch.setattr(lib, "vindex_bin", vindex_bin) + monkeypatch.setattr(lib, "call", call) + + +def _signature(hits): + """What a caller can see of a result: the RID text, the record's id, the score's type and its bits, in order.""" + return [ + (r.get_rid(), r.get("id"), type(score), struct.pack(" k and parts <= BUCKETS + assert [c["cap"] for c in calls.binary] == [k, k * parts] + assert calls.text == [] + assert _signature(again) == expected + assert shim._vindex_parts == parts + pq_index.find_nearest_approximate(pq_index.vectors[9], k=k) + assert [c["cap"] for c in calls.binary] == [ + k, + k * parts, + k * parts, + ] # one crossing now, with the size that was found + + +# queries the binary entry does not carry: the engine decides what they mean, so the JSON route must still be the one that runs, with +# the outcome (hits or error) the same as with the entry switched off +NOT_CARRIED = { + "a NaN": lambda v: np.concatenate([[np.nan], v[1:].astype(np.float64)]), + "an infinity": lambda v: np.concatenate([[np.inf], v[1:].astype(np.float64)]), + "beyond float32": lambda v: np.concatenate([[1e39], v[1:].astype(np.float64)]), + "an int array": lambda v: np.arange(DIM), + "numpy scalars": lambda v: [np.float32(x) for x in v], + "a None entry": lambda v: [None] + v[1:].tolist(), +} + + +def _outcome(call): + try: + return ("hits", _signature(call())) + except ArcadeDBError as e: + return ("error", str(e)) + + +@pytest.mark.parametrize("shape", sorted(NOT_CARRIED)) +def test_a_query_the_entry_cannot_carry_keeps_the_json_route( + pq_index, monkeypatch, shape +): + query = NOT_CARRIED[shape](pq_index.vectors[4]) + + def search(): + return _outcome(lambda: pq_index.find_nearest_approximate(query, k=5)) + + outcome, via_bin = _run(monkeypatch, search) + expected, _ = _run(monkeypatch, search, binary=False) + assert outcome == expected + assert via_bin.binary == [] and len(via_bin.text) == 1 + + +def test_a_wrong_dimension_is_the_engines_answer_on_both_routes(pq_index, monkeypatch): + """Thirteen float32 numbers cross by address; what the engine makes of a query of the wrong length is the same either way.""" + query = pq_index.vectors[4][: DIM - 3] + + def search(): + return _outcome(lambda: pq_index.find_nearest_approximate(query, k=5)) + + outcome, via_bin = _run(monkeypatch, search) + expected, _ = _run(monkeypatch, search, binary=False) + assert outcome == expected + assert len(via_bin.binary) == 1 and via_bin.binary[0]["dims"] == DIM - 3 + assert via_bin.text == [] + + +def test_a_filter_that_names_no_record_of_the_index_is_answered_alike( + pq_index, monkeypatch +): + """A plain RID of a bucket the type does not have: the engine's answer (hits or an error) is the same on both routes.""" + + def search(): + return _outcome( + lambda: pq_index.find_nearest_approximate( + pq_index.vectors[4], k=5, allowed_rids=["#999:0"] + ) + ) + + outcome, via_bin = _run(monkeypatch, search) + expected, _ = _run(monkeypatch, search, binary=False) + assert outcome == expected + assert len(via_bin.binary) == 1 and via_bin.binary[0]["pairs"] == 1 + + +def test_a_search_on_an_unindexed_property_raises_the_same_error(pq_index, monkeypatch): + """The library looks the index up by type and property on both routes: the same refusal comes back.""" + from arcadedb._native.schema_objects import NativeVectorIndex, _IndexShim + + bogus = NativeVectorIndex.__new__(NativeVectorIndex) + bogus._database = pq_index.database + bogus._shim = _IndexShim(pq_index.database, "Doc", "id", "Doc[id]") + monkeypatch.setattr(bogus, "get_quantization", lambda: "PRODUCT") + + def search(): + return _outcome( + lambda: bogus.find_nearest_approximate(pq_index.vectors[4], k=5) + ) + + outcome, via_bin = _run(monkeypatch, search) + expected, _ = _run(monkeypatch, search, binary=False) + assert outcome[0] == "error" and "No vector index on Doc[id]" in outcome[1] + assert outcome == expected + assert len(via_bin.binary) == 1 diff --git a/tests/test_native_vindex_glue.py b/tests/test_native_vindex_glue.py new file mode 100644 index 0000000..2a0b35b --- /dev/null +++ b/tests/test_native_vindex_glue.py @@ -0,0 +1,425 @@ +"""The Python half of the binary approximate search of the native wheel: which queries cross by address, what the library is handed, and +how its answer is read (``arcadedb_n_vindex_bin``, ``find_nearest_approximate``). + +Before, the query vector crossed as a JSON list of d numbers (``arcadedb_n_vindex``, op ``approximate``) and the hits came back as a JSON +list of ``[rid, score]``. Now a query that float32 carries goes by address (a float32 array and its count), a filter goes as int64 +(bucket, position) pairs, and the library writes the hits into a buffer of int64 words that Python hands it: three words per hit, the +bucket, the position and the score as the bits of a double. Anything the entry cannot carry keeps the JSON route. + +Needs no library: ``lib.vindex_bin`` is replaced by a stand-in that reads the buffers by their addresses, as the library does, so the +tests see exactly the floats and the pairs the engine would be given. The real library, and both routes against each other, are in +``test_native_vindex_bin.py`` (native backend only). +""" + +import ctypes +import json +import struct + +import numpy as np +import pytest + +from arcadedb._native import lib, schema_objects +from arcadedb._native.schema_objects import NativeVectorIndex, _IndexShim +from arcadedb.exceptions import ArcadeDBError + + +def _f32(x): + """The float32 that a C cast of the double `x` gives (what Java's floatValue() gives), by way of struct, not numpy.""" + return struct.unpack(" cap: + return len(hits) + words = (ctypes.c_int64 * (3 * len(hits))).from_address(out) + for i, (bucket, position, score) in enumerate(hits): + words[3 * i] = bucket + words[3 * i + 1] = position + words[3 * i + 2] = struct.unpack("