From a1200816fdd82e2d70a6c60eb1133389f452327c Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Tue, 3 Nov 2026 08:23:40 +0900 Subject: [PATCH 1/3] Native find_nearest: send the query vector as one float32 array, not a Python list NativeVectorIndex._run_search turned every query into a list with tolist(). The codec then checked the list for float32-exactness (_exact_float_list: a pass of numpy over the list, tag BL) or, for a float64 query whose values float32 rounds, printed JSON number text (about 611 us at 1536 dimensions, ArcadeData/arcadedb-python#107). A numeric 1-d query (an ndarray of any numeric dtype, or a list/tuple of numbers, which numpy walks in C) now goes to the codec as a float32 array and travels through the float32 side buffer (tag B) directly. The engine reads the vector as float32 whatever it is given, so the rounding is unchanged. Anything else (text or None entries, a generator, ragged input, numpy missing) stays a list and reaches the engine as before. Port of humemai/arcadedb-embedded-python#523 (backlog row N-07) onto main after #107: the per-element _prep calls it removed are gone since #107, what remains is tolist() plus the exactness check, and the float64 text path. Untimed; a timed A/B decides it. Tests: test_find_nearest_query_forms.py (both backends: every form answers the same neighbours at the same distances, refusals unchanged; passes before and after) and test_native_find_nearest_params.py (native: where the floats travel and that the list exactness check is not run; 5 of its 6 tests fail on the code before). Co-Authored-By: Claude Opus 5.5 --- src/arcadedb/_native/schema_objects.py | 30 +++++-- tests/conftest.py | 2 + tests/test_find_nearest_query_forms.py | 101 +++++++++++++++++++++ tests/test_native_find_nearest_params.py | 110 +++++++++++++++++++++++ 4 files changed, 238 insertions(+), 5 deletions(-) create mode 100644 tests/test_find_nearest_query_forms.py create mode 100644 tests/test_native_find_nearest_params.py diff --git a/src/arcadedb/_native/schema_objects.py b/src/arcadedb/_native/schema_objects.py index c1ccbfe..a7a011d 100644 --- a/src/arcadedb/_native/schema_objects.py +++ b/src/arcadedb/_native/schema_objects.py @@ -9,6 +9,7 @@ from ..exceptions import ArcadeDBError from ..vector import VectorIndex as _BaseVectorIndex from . import needs_jvm +from .codec import _np _PROPERTY = re.compile(r"name: (.+?), type: ([A-Z_]+)") @@ -169,6 +170,29 @@ def _rid_text(value) -> str: return get_rid() if get_rid is not None else str(value) +def _query_vector_param(query_vector): + """The query vector as the one parameter the library takes it in. + + A numeric 1-d vector (an ndarray, or a list/tuple of numbers: numpy walks a list in C) becomes a float32 array, which + ``encode_params`` sends through the binary float32 side buffer (tag ``B``): no ``tolist()``, no float32-exactness check of a + list, and no JSON number text for a float64 query whose values float32 would round. The engine reads the vector as float32 + whatever it is given, so the rounding is the same. Anything else (text or None entries, a generator, numpy missing) stays a + list and reaches the engine as it was given; numpy would turn a None into nan. + """ + if _np is not None: + try: + arr = _np.asarray(query_vector) + except ( + Exception + ): # nosec B110 - ragged or unconvertible input: the list path below reports it + arr = None + if arr is not None and arr.ndim == 1 and arr.dtype.kind in "fiu": + return arr.astype(_np.float32, copy=False) + return ( + query_vector.tolist() if hasattr(query_vector, "tolist") else list(query_vector) + ) + + class NativeVectorIndex(_BaseVectorIndex): """``VectorIndex`` over ``vectorNeighbors``: ``find_nearest`` returns ``[(record, distance)]`` like the JPype wrapper.""" @@ -192,11 +216,7 @@ def _run_search( return self._run_approximate(query_vector, k, allowed_rids) from .records import _wrap_dict - vector = ( - query_vector.tolist() - if hasattr(query_vector, "tolist") - else list(query_vector) - ) + vector = _query_vector_param(query_vector) # the statement's own options: efSearch, and filter (the RIDs the search may return), as the JPype path passes them options = {} if ef_search is not None: diff --git a/tests/conftest.py b/tests/conftest.py index c5961c2..b6f0b3c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -38,6 +38,8 @@ if BACKEND != "native": # Java -> Python upcalls (Python functions in SQL) exist on the native build only collect_ignore.append("test_native_udf.py") + # how find_nearest hands the query vector to the native library (the codec's side buffer) + collect_ignore.append("test_native_find_nearest_params.py") # the columnar insert's stored records are compared between the JSON and the columnar entry points of the native library collect_ignore.append("test_native_columnar_insert.py") # the native build's ingest defaults (insert_many and graph_batch commit size), observed on the native library diff --git a/tests/test_find_nearest_query_forms.py b/tests/test_find_nearest_query_forms.py new file mode 100644 index 0000000..a9e1d16 --- /dev/null +++ b/tests/test_find_nearest_query_forms.py @@ -0,0 +1,101 @@ +"""The forms a ``find_nearest`` query can take, on both backends: every numeric form answers the same neighbours at the same +distances, and a query that is not a numeric vector of the index's dimension is still refused. + +The native backend hands the library a float32 array for a numeric query (``test_native_find_nearest_params.py`` pins how it +travels); these tests are the contract that change must keep, and they pass unchanged on the code before it. +""" + +import arcadedb +import numpy as np +import pytest +from arcadedb import ArcadeDBError + +DIM = 768 +N = 60 + + +@pytest.fixture +def index(temp_db_path): + db = arcadedb.create_database(temp_db_path) + db.command("sql", "CREATE VERTEX TYPE Doc") + db.command("sql", "CREATE PROPERTY Doc.id INTEGER") + db.command("sql", "CREATE PROPERTY Doc.embedding ARRAY_OF_FLOATS") + vectors = np.random.default_rng(7).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] + ) + idx = db.create_vector_index("Doc", "embedding", dimensions=DIM) + idx.vectors = vectors + yield idx + db.drop() + + +def _answer(hits): + return [(record.get("id"), distance) for record, distance in hits] + + +def test_every_form_of_the_query_returns_the_same_neighbours_at_the_same_distances( + index, +): + query = index.vectors[11] + expected = _answer(index.find_nearest(query.tolist(), k=10)) + assert len(expected) == 10 + assert expected[0][0] == 11 + forms = { + "float32 array": query, + "float64 array": query.astype(np.float64), + "tuple": tuple(query.tolist()), + "list of numpy floats": list(query), + "generator": (float(x) for x in query), + "non-contiguous view": np.repeat(query, 2)[::2], + } + for name, form in forms.items(): + assert _answer(index.find_nearest(form, k=10)) == expected, name + + +def test_a_float64_query_that_float32_would_round_answers_like_its_float32_form( + index, +): + # not float32-exact: the values float32 rounds are the ones that left the native list path as JSON text + query = np.random.default_rng(3).standard_normal(DIM) + assert not np.array_equal(query.astype(np.float32).astype(np.float64), query) + expected = _answer(index.find_nearest(query.astype(np.float32), k=10)) + assert _answer(index.find_nearest(query, k=10)) == expected + assert _answer(index.find_nearest(query.tolist(), k=10)) == expected + + +def test_an_integer_array_is_a_numeric_query(index): + ints = np.arange(DIM) % 5 + expected = _answer(index.find_nearest(ints.tolist(), k=5)) + assert _answer(index.find_nearest(ints, k=5)) == expected + + +def test_find_nearest_by_key_agrees_with_the_array_query(index): + by_key = _answer(index.find_nearest_by_key(4, k=6)) + assert by_key == _answer(index.find_nearest(index.vectors[4], k=6)) + + +def test_filter_and_ef_search_options_still_apply_to_an_array_query(index): + everything = index.find_nearest(index.vectors[2], k=10) + allowed = [record.get_rid() for record, _ in everything[:3]] + hits = index.find_nearest(index.vectors[2], k=10, allowed_rids=allowed) + assert {record.get_rid() for record, _ in hits} == set(allowed) + assert _answer(index.find_nearest(index.vectors[2], k=10, ef_search=64)) + + +def test_a_query_of_the_wrong_dimension_is_refused_in_both_forms(index): + with pytest.raises(ArcadeDBError): + index.find_nearest(np.zeros(5, dtype=np.float32), k=3) + with pytest.raises(ArcadeDBError): + index.find_nearest([0.0] * 5, k=3) + + +@pytest.mark.parametrize( + "bad", [[None] * DIM, ["x"] * DIM], ids=["None entries", "text entries"] +) +def test_a_query_that_is_not_numeric_is_refused(index, bad): + # numpy would turn None into nan; the query must reach the engine as it was given + with pytest.raises(ArcadeDBError): + index.find_nearest(bad, k=3) diff --git a/tests/test_native_find_nearest_params.py b/tests/test_native_find_nearest_params.py new file mode 100644 index 0000000..db94a95 --- /dev/null +++ b/tests/test_native_find_nearest_params.py @@ -0,0 +1,110 @@ +"""``find_nearest`` hands the native library its query vector as one float32 array, not as a Python list of floats. + +``find_nearest`` used to turn every query into a list with ``tolist()``. The codec then checked the list for float32-exactness +(a pass of numpy over the list, ``_exact_float_list``) and sent it in the side buffer as tag ``BL``, or, for a float64 query whose +values float32 rounds, printed one float repr per element as JSON text (611 us at 1536 dimensions). A numpy array travels through +the float32 side buffer directly (tag ``B``: a pointer and a count, read by the library straight from the C buffer). Whether +that is faster is a timed A/B, not these tests; they count what the codec does and where the floats travel. + +The answers are pinned for both backends in ``test_find_nearest_query_forms.py``. This file needs the native library, so +conftest.py does not collect it on the JVM backend. +""" + +import arcadedb +import numpy as np +import pytest +from arcadedb._native import codec +from arcadedb._native import database as native_database + +DIM = 768 +N = 40 + + +@pytest.fixture +def index(temp_db_path): + db = arcadedb.create_database(temp_db_path) + db.command("sql", "CREATE VERTEX TYPE Doc") + db.command("sql", "CREATE PROPERTY Doc.id INTEGER") + db.command("sql", "CREATE PROPERTY Doc.embedding ARRAY_OF_FLOATS") + vectors = np.random.default_rng(7).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] + ) + idx = db.create_vector_index("Doc", "embedding", dimensions=DIM) + idx.vectors = vectors + yield idx + db.drop() + + +def _search_encodings(monkeypatch): + """Records what ``encode_params`` returns for every statement; the search is the one carrying DIM floats beside its text.""" + seen = [] + real = native_database.encode_params + + def recording(args, binary=False): + encoded = real(args, binary) + seen.append(encoded) + return encoded + + monkeypatch.setattr(native_database, "encode_params", recording) + return seen + + +def _searches(seen): + return [text for text, _, floats in seen if floats == DIM] + + +def test_an_array_query_travels_in_the_float32_side_buffer(index, monkeypatch): + seen = _search_encodings(monkeypatch) + index.find_nearest(index.vectors[3], k=5) + search = _searches(seen) + assert len(search) == 1 + assert '"B"' in search[0] + assert '"BL"' not in search[0] + + +@pytest.mark.parametrize( + "form", + ["float64 array", "float64 list", "float32 list"], +) +def test_every_numeric_form_travels_in_the_side_buffer_as_an_array( + index, monkeypatch, form +): + exact = index.vectors[3] + rounding = np.random.default_rng(5).standard_normal(DIM) # float32 would round it + query = { + "float64 array": rounding, + "float64 list": rounding.tolist(), + "float32 list": exact.tolist(), + }[form] + seen = _search_encodings(monkeypatch) + index.find_nearest(query, k=5) + search = _searches(seen) + # before: the float64 forms were JSON text (no side buffer at all), the float32 list tag BL + assert len(search) == 1 + assert '"B"' in search[0] + assert '"BL"' not in search[0] + + +def test_a_numeric_query_is_not_walked_for_float32_exactness(index, monkeypatch): + calls = [] + real = codec._exact_float_list + + def counting(value): + calls.append(len(value)) + return real(value) + + monkeypatch.setattr(codec, "_exact_float_list", counting) + index.find_nearest(index.vectors[3], k=5) + index.find_nearest(index.vectors[3].tolist(), k=5) + index.find_nearest(index.vectors[3].astype(np.float64), k=5) + assert [n for n in calls if n >= DIM] == [] + + +def test_a_query_that_is_not_numeric_stays_a_list(index, monkeypatch): + seen = _search_encodings(monkeypatch) + with pytest.raises(arcadedb.ArcadeDBError): + index.find_nearest([None] * DIM, k=3) + assert _searches(seen) == [] From 0e2fb51ff55d540c634f8d3d471dc9eae3704160 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Tue, 3 Nov 2026 08:23:57 +0900 Subject: [PATCH 2/3] tests: isort order for the find_nearest query-form tests Co-Authored-By: Claude Opus 5.5 --- tests/test_find_nearest_query_forms.py | 3 ++- tests/test_native_find_nearest_params.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_find_nearest_query_forms.py b/tests/test_find_nearest_query_forms.py index a9e1d16..848ba14 100644 --- a/tests/test_find_nearest_query_forms.py +++ b/tests/test_find_nearest_query_forms.py @@ -5,9 +5,10 @@ travels); these tests are the contract that change must keep, and they pass unchanged on the code before it. """ -import arcadedb import numpy as np import pytest + +import arcadedb from arcadedb import ArcadeDBError DIM = 768 diff --git a/tests/test_native_find_nearest_params.py b/tests/test_native_find_nearest_params.py index db94a95..6b2f7c5 100644 --- a/tests/test_native_find_nearest_params.py +++ b/tests/test_native_find_nearest_params.py @@ -10,9 +10,10 @@ conftest.py does not collect it on the JVM backend. """ -import arcadedb import numpy as np import pytest + +import arcadedb from arcadedb._native import codec from arcadedb._native import database as native_database From c91a8287d689991d6b62866f62e3c8382307cb27 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Wed, 4 Nov 2026 11:03:27 +0900 Subject: [PATCH 3/3] Native find_nearest: include_vector=False returns hits without the indexed vector, loaded by RID on first read find_nearest keeps its default: every hit carries its vector, and the default search is the code of the query-as-float32 change (#203), untouched. A new keyword-only include_vector=False asks for hits built without the vector, for callers that need only ids or distances. On the native wheel the search then runs SELECT *, !, !record FROM (SELECT expand(vectorNeighbors(...))) and each hit loads its record once by RID the first time the vector or the whole record is needed. The JVM build accepts the keyword and ignores it (it loads whole records either way). Reworks the earlier lazy-by-default version after its timed A/B: hits not read 0.66 and 0.805 (1536 and 256 dims), but vector read 1.011 and 1.336, so lazy-by-default is a cost for callers that read the vector. Tests run on both backends (values equal lookup_by_rid, keyword accepted) and on the native wheel (default ships the vector with no lookup and no dropped column, opt-in ships none and looks up once per record). Co-Authored-By: Claude Opus 5.5 --- docs/modules/ROOT/pages/how-to/vectors.adoc | 18 +- src/arcadedb/_native/records.py | 28 +++ src/arcadedb/_native/schema_objects.py | 73 +++++++- src/arcadedb/vector.py | 23 ++- tests/conftest.py | 2 + tests/test_native_vector_hit_lazy.py | 133 ++++++++++++++ tests/test_vector_hit_records.py | 186 ++++++++++++++++++++ 7 files changed, 459 insertions(+), 4 deletions(-) create mode 100644 tests/test_native_vector_hit_lazy.py create mode 100644 tests/test_vector_hit_records.py diff --git a/docs/modules/ROOT/pages/how-to/vectors.adoc b/docs/modules/ROOT/pages/how-to/vectors.adoc index db9cc54..f1124dc 100644 --- a/docs/modules/ROOT/pages/how-to/vectors.adoc +++ b/docs/modules/ROOT/pages/how-to/vectors.adoc @@ -157,7 +157,7 @@ print([row["name"] for row in rows]) === find_nearest -`VectorIndex.find_nearest(query_vector, k=10, ef_search=None, allowed_rids=None)` searches without a statement. It takes a list or a NumPy array and returns a list of `(record, distance)` tuples, nearest first. Use it when you already hold the index object, or when you have a list of record ids to search within. +`VectorIndex.find_nearest(query_vector, k=10, ef_search=None, allowed_rids=None, *, include_vector=True)` searches without a statement. It takes a list or a NumPy array and returns a list of `(record, distance)` tuples, nearest first. Use it when you already hold the index object, or when you have a list of record ids to search within. [source,python] ---- @@ -171,6 +171,22 @@ python-ish python ---- +=== Hits without the vector + +Every hit of `find_nearest` carries its indexed vector, so a search at 1536 dimensions ships and decodes 1536 numbers per hit whether you read them or not. If you need only the ids, the distances or other properties of the hits, pass `include_vector=False`. On the native build each hit is then built without the vector. It loads the vector, with the rest of the record, in one lookup by record id the first time you read it (`get` of the vector, `has_property`, `to_dict`, `get_property_names`, `modify` or `save`), so a loop that reads the vector of every hit pays one lookup per hit and does more work than the default. The values are the ones `db.lookup_by_rid` returns, and a hit deleted since the search raises `ArcadeDBError` on that first read. The arcadedb-jvm package accepts the argument and returns whole records either way. + +[source,python] +---- +for record, distance in index.find_nearest(query, k=2, include_vector=False): + print(record.get("name")) +---- + +[source,text] +---- +python-ish +python +---- + === Filtered search `allowed_rids` restricts `find_nearest` to the records whose ids you list, as strings. The search considers only those records, so it still returns `k` of them. diff --git a/src/arcadedb/_native/records.py b/src/arcadedb/_native/records.py index f790ade..0e64e71 100644 --- a/src/arcadedb/_native/records.py +++ b/src/arcadedb/_native/records.py @@ -50,9 +50,17 @@ def __init__(self, database, type_name, rid=None, props=None, mutable=False): self._props: Dict[str, Any] = dict(props or {}) self._mutable = mutable or not rid self._property_names_cache = None + # names of properties a search hit was built without (see ``_wrap_partial``); None for every other record + self._lazy = None # -- reading + def _fill(self) -> None: + """Load the properties a search hit was built without, in one lookup by RID (engine order); nothing to do for any other record.""" + if self._lazy: + self._props = dict(self._db.lookup_by_rid(self._rid)._props) + self._lazy = None + def _check_open(self) -> None: database = self._database if database is not None and database._closed and self._rid: @@ -63,6 +71,8 @@ def _check_open(self) -> None: def get(self, name: str, convert_types: bool = True) -> Any: self._check_open() + if self._lazy and name in self._lazy: + self._fill() return self._props.get(name) def get_raw(self, name: str) -> Any: @@ -70,14 +80,18 @@ def get_raw(self, name: str) -> Any: def has_property(self, name: str) -> bool: self._check_open() + if self._lazy and name in self._lazy: + self._fill() return name in self._props def get_property_names(self) -> List[str]: self._check_open() + self._fill() return list(self._props) def to_dict(self, convert_types: bool = True) -> Dict[str, Any]: self._check_open() + self._fill() return dict(self._props) def get_rid(self) -> str: @@ -106,6 +120,7 @@ def set(self, name: str, value: Any): return self def save(self): + self._fill() db = self._db if self._rid is None: from .extras import veto_count @@ -138,6 +153,7 @@ def delete(self): def modify(self): self._check_open() + self._fill() return type(self)(self._db, self._type, self._rid, self._props, mutable=True) def __repr__(self) -> str: @@ -218,3 +234,15 @@ def _wrap_dict(database, record: dict, kind: str = "d"): type_name = props.pop("@type", None) kind = props.pop("@cat", kind) return _CLASS_OF.get(kind, Document)(database, type_name, rid, props) + + +def _wrap_partial(database, rid, type_name, kind, props, omitted): + """The wrapper for a search hit built without the properties named in ``omitted`` (the indexed vector: not shipped with every hit). + + Reading one of them, or anything that needs the whole record (``to_dict``, ``get_property_names``, ``modify``, ``save``), loads the + record once by RID; ``get`` and ``has_property`` of any other property are answered from what was shipped. + """ + cls = {"v": Vertex, "e": Edge}.get(kind, Document) + record = cls(database, type_name, rid, props) + record._lazy = frozenset(omitted) + return record diff --git a/src/arcadedb/_native/schema_objects.py b/src/arcadedb/_native/schema_objects.py index a7a011d..5f69e0d 100644 --- a/src/arcadedb/_native/schema_objects.py +++ b/src/arcadedb/_native/schema_objects.py @@ -193,6 +193,18 @@ def _query_vector_param(query_vector): ) +def _quote_name(name: str) -> str: + return "`" + name.replace("`", "``") + "`" + + +# find_nearest(include_vector=False): the engine's row of a hit is the record's properties plus the columns @rid, @type and distance +# (ScoredRecordRow; before 26.11.1, ArcadeData/arcadedb#10898, also a nested `record` column holding every property a second time). +# `*` copies the row, `!name` drops the indexed vector and the `record` column. A property of the record called distance or record is +# hidden by those columns in the row, so such a hit holds them back too (_LAZY_COLUMNS) and loads them with the vector. +_LAZY_COLUMNS = ("record", "distance") +_HITS_FROM = "FROM (SELECT expand(vectorNeighbors(?, ?, ?)))" + + class NativeVectorIndex(_BaseVectorIndex): """``VectorIndex`` over ``vectorNeighbors``: ``find_nearest`` returns ``[(record, distance)]`` like the JPype wrapper.""" @@ -210,10 +222,21 @@ def _check_live(self): return def _run_search( - self, *, query_vector, k, allowed_rids, approximate, ef_search=None + self, + *, + query_vector, + k, + allowed_rids, + approximate, + ef_search=None, + include_vector=True, ): if approximate: return self._run_approximate(query_vector, k, allowed_rids) + if not include_vector: + return self._run_search_without_vector( + query_vector, k, allowed_rids, ef_search + ) from .records import _wrap_dict vector = _query_vector_param(query_vector) @@ -257,6 +280,54 @@ def _run_search( out.append((record, float(row.get("distance")))) return self._sort_results(out)[:k] + def _run_search_without_vector(self, query_vector, k, allowed_rids, ef_search): + """``find_nearest(include_vector=False)``: the hits come back without the indexed vector, which loads on first use. + + ``expand(vectorNeighbors(..))`` alone ships every hit's vector, and the engine reads each record to build the row whatever + the columns asked for. Here the statement drops the vector (and the nested ``record`` column of an engine before 26.11.1), + and ``_Mixin._fill`` loads the record, once, by RID, the first time the vector or the whole record is needed. + """ + from .records import _wrap_partial + + vector = _query_vector_param(query_vector) + options = {} + if ef_search is not None: + options["efSearch"] = int(ef_search) + if allowed_rids: + options["filter"] = [_rid_text(r) for r in allowed_rids] + prop = self._shim._prop + dropped = ", ".join( + "!" + _quote_name(name) for name in dict.fromkeys((prop, "record")) + ) + sql = f"SELECT *, {dropped} {_HITS_FROM}" # nosec B608 - dropped is backtick-quoted schema names, the values go through ? parameters + if options: + sql = sql.replace("vectorNeighbors(?, ?, ?)", "vectorNeighbors(?, ?, ?, ?)") + rows = self._database.query( + "sql", sql, self._shim._name, vector, int(k), options + ) + else: + rows = self._database.query("sql", sql, self._shim._name, vector, int(k)) + omitted = (prop, *_LAZY_COLUMNS) + out = [] + for row in rows: + props = row.to_dict() + rid = props.pop("@rid", None) + type_name = props.pop("@type", None) + distance = props.pop("distance") + if rid is None: + record = None + else: + record = _wrap_partial( + self._database, + _rid_text(rid), + type_name, + self._database._kind_of(type_name), + props, + omitted, + ) + out.append((record, float(distance))) + return self._sort_results(out)[:k] + 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 diff --git a/src/arcadedb/vector.py b/src/arcadedb/vector.py index f2c69dc..96f3ca4 100644 --- a/src/arcadedb/vector.py +++ b/src/arcadedb/vector.py @@ -566,6 +566,7 @@ def _run_search( allowed_rids, approximate, ef_search=None, + include_vector=True, # accepted for API parity, the JVM build loads whole records ): java_vector = to_java_float_array(query_vector) allowed_rids_set = self._build_allowed_rids_set(allowed_rids) @@ -584,6 +585,8 @@ def find_nearest( k=10, ef_search=None, allowed_rids=None, + *, + include_vector=True, ): """ Find k nearest neighbors to the query vector. @@ -604,6 +607,17 @@ def find_nearest( uses the Java index default/adaptive behavior. allowed_rids: Optional list of RID strings (e.g. ["#1:0", "#2:5"]) to restrict search + include_vector: Keyword only, default ``True``: every hit carries its + indexed vector, as always. Pass ``False`` when you need only the + ids or the distances (or other properties) of the hits. On the + native build each hit is then built without the indexed vector, and + loads it, with the rest of the record, in one lookup by RID the + first time you read it (``get`` of the vector, ``has_property``, + ``to_dict``, ``get_property_names``, ``modify``, ``save``); the + values are those of ``lookup_by_rid``, and a hit deleted since the + search raises ``ArcadeDBError`` on that first read. The JVM build + accepts the argument and ignores it: it loads whole records either + way, so the call costs the same. Returns: List of tuples: [(record, score), ...] @@ -611,9 +625,13 @@ def find_nearest( - score: The similarity score/distance """ self._check_live() - return self._find_nearest_live(query_vector, k, ef_search, allowed_rids) + return self._find_nearest_live( + query_vector, k, ef_search, allowed_rids, include_vector + ) - def _find_nearest_live(self, query_vector, k, ef_search, allowed_rids): + def _find_nearest_live( + self, query_vector, k, ef_search, allowed_rids, include_vector=True + ): # find_nearest after _check_live(): find_nearest_by_key checks once and # comes here, so a call never pays the isValid() crossing twice. effective_ef_search = self._normalize_ef_search(ef_search) @@ -625,6 +643,7 @@ def _find_nearest_live(self, query_vector, k, ef_search, allowed_rids): allowed_rids=allowed_rids, approximate=False, ef_search=effective_ef_search, + include_vector=include_vector, ) except ArcadeDBError: diff --git a/tests/conftest.py b/tests/conftest.py index cc4c0d3..0931a92 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -40,6 +40,8 @@ collect_ignore.append("test_native_udf.py") # how find_nearest hands the query vector to the native library (the codec's side buffer) collect_ignore.append("test_native_find_nearest_params.py") + # find_nearest(include_vector=False) on the native wheel fills a hit's vector on first use (counts lookups of the native wrapper) + collect_ignore.append("test_native_vector_hit_lazy.py") # the columnar insert's stored records are compared between the JSON and the columnar entry points of the native library collect_ignore.append("test_native_columnar_insert.py") # the native build's ingest defaults (insert_many and graph_batch commit size), observed on the native library diff --git a/tests/test_native_vector_hit_lazy.py b/tests/test_native_vector_hit_lazy.py new file mode 100644 index 0000000..7b672f7 --- /dev/null +++ b/tests/test_native_vector_hit_lazy.py @@ -0,0 +1,133 @@ +"""Native wheel: ``find_nearest`` keeps its default (hits carry the vector, no extra lookup) and, asked for ``include_vector=False``, +builds a hit without the indexed vector and fills it on first use, with one lookup by RID. + +The values themselves are checked on both wheels in ``test_vector_hit_records.py``; this file counts what is shipped and looked up, +and which statement the search runs. +""" + +import numpy as np +import pytest + +from arcadedb.exceptions import ArcadeDBError + +DIM = 8 + + +@pytest.fixture +def lazy_db(temp_db, monkeypatch): + temp_db.command("sql", "CREATE VERTEX TYPE Doc BUCKETS 4") + temp_db.command("sql", "CREATE PROPERTY Doc.id INTEGER") + temp_db.command("sql", "CREATE PROPERTY Doc.embedding ARRAY_OF_FLOATS") + rng = np.random.default_rng(5) + vecs = rng.standard_normal((40, DIM)).astype(np.float32) + with temp_db.transaction(): + for i in range(40): + temp_db.command( + "sql", + "INSERT INTO Doc SET id = ?, name = ?, embedding = ?", + i, + f"n{i}", + vecs[i].tolist(), + ) + index = temp_db.create_vector_index( + "Doc", "embedding", dimensions=DIM, id_property="id" + ) + lookups = [] + original = type(temp_db).lookup_by_rid + + def counting(self, rid): + lookups.append(str(rid)) + return original(self, rid) + + monkeypatch.setattr(type(temp_db), "lookup_by_rid", counting) + statements = [] + original_query = type(temp_db).query + + def recording(self, language, text, *args, **kwargs): + statements.append(text) + return original_query(self, language, text, *args, **kwargs) + + monkeypatch.setattr(type(temp_db), "query", recording) + return temp_db, index, vecs, lookups, statements + + +def _search_statements(statements): + return [s for s in statements if "vectorNeighbors" in s] + + +def test_the_default_search_ships_the_vector_and_looks_up_nothing(lazy_db): + db, index, vecs, lookups, statements = lazy_db + hits = index.find_nearest(vecs[2], k=6) + assert len(hits) == 6 + for record, _ in hits: + assert "embedding" in record._props + assert record._lazy is None + assert len(record.get("embedding")) == DIM + assert record.has_property("embedding") + assert "embedding" in record.to_dict() + assert lookups == [] + (text,) = _search_statements(statements) + assert "!" not in text and "SELECT *" not in text + + +def test_asking_for_the_default_in_words_is_the_default(lazy_db): + db, index, vecs, lookups, statements = lazy_db + hits = index.find_nearest(vecs[2], k=3, include_vector=True) + assert all("embedding" in record._props for record, _ in hits) + assert lookups == [] + (text,) = _search_statements(statements) + assert "!" not in text + + +def test_an_opt_in_search_ships_no_vector_and_looks_up_nothing(lazy_db): + db, index, vecs, lookups, statements = lazy_db + hits = index.find_nearest(vecs[2], k=6, include_vector=False) + assert len(hits) == 6 + for record, _ in hits: + assert "embedding" not in record._props + assert record.get("id") is not None and record.get("name") + assert record.get_rid() and record.get_type_name() == "Doc" + assert lookups == [] + (text,) = _search_statements(statements) + assert "!`embedding`" in text + + +def test_the_vector_is_one_lookup_per_record_however_it_is_reached(lazy_db): + db, index, vecs, lookups, statements = lazy_db + hits = index.find_nearest(vecs[2], k=4, include_vector=False) + first, second, third, fourth = (r for r, _ in hits) + assert len(first.get("embedding")) == DIM + assert len(first.get("embedding")) == DIM + assert first.has_property("embedding") + assert "embedding" in first.to_dict() + assert "embedding" in first.get_property_names() + assert len(lookups) == 1 + assert second.has_property("embedding") + assert len(lookups) == 2 + assert "embedding" in third.to_dict() + assert "embedding" in fourth.get_property_names() + assert len(lookups) == 4 + assert sorted(lookups) == sorted(r.get_rid() for r, _ in hits) + + +def test_the_loaded_vector_is_the_one_the_default_search_ships(lazy_db): + db, index, vecs, lookups, statements = lazy_db + full = { + r.get_rid(): list(r.get("embedding")) + for r, _ in index.find_nearest(vecs[9], k=5) + } + assert lookups == [] + lean = index.find_nearest(vecs[9], k=5, include_vector=False) + assert lookups == [] + for record, _ in lean: + assert list(record.get("embedding")) == full[record.get_rid()] + assert len(lookups) == 5 + + +def test_the_vector_of_a_closed_database_is_not_looked_up(lazy_db): + db, index, vecs, lookups, statements = lazy_db + record, _ = index.find_nearest(vecs[2], k=1, include_vector=False)[0] + db.close() + with pytest.raises(ArcadeDBError, match="closed"): + record.get("embedding") + assert lookups == [] diff --git a/tests/test_vector_hit_records.py b/tests/test_vector_hit_records.py new file mode 100644 index 0000000..f8774fc --- /dev/null +++ b/tests/test_vector_hit_records.py @@ -0,0 +1,186 @@ +"""The records a ``find_nearest`` hit holds, on both backends: the same values as a lookup by RID, whichever way it is asked. + +``find_nearest(..., include_vector=False)`` builds a hit without the indexed vector on the native wheel and fills it on first +use (``test_native_vector_hit_lazy.py`` counts that); the JVM wheel accepts the keyword and returns whole records. These tests are +the contract both builds keep: type, property names and their order, every value, the vector included, and a record that can be +modified and saved after a search. The default call (``include_vector=True``) is held to the same contract. +""" + +import numpy as np +import pytest + +DIM = 8 +N_DOC = 48 +N_SUB = 16 +BOTH = pytest.mark.parametrize("include_vector", [True, False]) + + +def _make(temp_db, extra_properties, extra_columns, extra_values): + temp_db.command("sql", "CREATE VERTEX TYPE Doc BUCKETS 4") + temp_db.command("sql", "CREATE VERTEX TYPE SubDoc EXTENDS Doc") + temp_db.command("sql", "CREATE PROPERTY Doc.id INTEGER") + temp_db.command("sql", "CREATE PROPERTY Doc.label STRING") + for statement in extra_properties: + temp_db.command("sql", statement) + temp_db.command("sql", "CREATE PROPERTY Doc.embedding ARRAY_OF_FLOATS") + rng = np.random.default_rng(11) + vecs = rng.standard_normal((N_DOC + N_SUB, DIM)).astype(np.float32) + with temp_db.transaction(): + for i in range(N_DOC + N_SUB): + type_name = "Doc" if i < N_DOC else "SubDoc" + temp_db.command( + "sql", + f"INSERT INTO {type_name} SET id = ?, label = ?{extra_columns}, embedding = ?", + i, + f"label-{i}", + *extra_values(i), + vecs[i].tolist(), + ) + index = temp_db.create_vector_index( + "Doc", "embedding", dimensions=DIM, id_property="id" + ) + return temp_db, index, vecs + + +@pytest.fixture +def hit_db(temp_db): + """Four buckets and a subtype.""" + return _make(temp_db, [], "", lambda i: ()) + + +@pytest.fixture +def clash_db(temp_db): + """As ``hit_db``, plus properties called ``distance`` and ``record`` (the names of the search row's own columns).""" + return _make( + temp_db, + [ + "CREATE PROPERTY Doc.distance DOUBLE", + "CREATE PROPERTY Doc.record STRING", + ], + ", distance = ?, record = ?", + lambda i: (1000.5 + i, f"record-{i}"), + ) + + +def _same_record(hit, ref): + assert hit.get_rid() == ref.get_rid() + assert hit.get_type_name() == ref.get_type_name() + assert type(hit).__name__ == type(ref).__name__ + assert list(hit.get_property_names()) == list(ref.get_property_names()) + assert list(hit.to_dict()) == list(ref.to_dict()) + for name in ref.get_property_names(): + assert hit.has_property(name) + assert hit.get(name) == ref.get(name), name + + +@BOTH +def test_hits_equal_a_lookup_by_rid(hit_db, include_vector): + db, index, vecs = hit_db + seen_types = set() + for q in (0, 5, 40, 50, 60): + hits = index.find_nearest(vecs[q], k=7, include_vector=include_vector) + assert len(hits) == 7 + assert [d for _, d in hits] == sorted(d for _, d in hits) + for record, score in hits: + ref = db.lookup_by_rid(record.get_rid()) + seen_types.add(record.get_type_name()) + _same_record(record, ref) + assert seen_types == {"Doc", "SubDoc"} + + +def test_the_keyword_changes_no_answer(hit_db): + db, index, vecs = hit_db + for q in (1, 20, 55): + full = index.find_nearest(vecs[q], k=9) + lean = index.find_nearest(vecs[q], k=9, include_vector=False) + assert [(r.get_rid(), d) for r, d in full] == [ + (r.get_rid(), d) for r, d in lean + ] + assert [r.get("id") for r, _ in full] == [r.get("id") for r, _ in lean] + for (a, _), (b, _) in zip(full, lean): + assert list(a.get("embedding")) == list(b.get("embedding")) + + +def test_the_keyword_is_keyword_only_and_a_boolean_default(hit_db): + db, index, vecs = hit_db + with pytest.raises(TypeError): + index.find_nearest(vecs[0], 3, None, None, False) + assert index.find_nearest(vecs[0], k=3, include_vector=True) + + +def test_hits_without_the_vector_equal_a_lookup_when_properties_are_named_like_columns( + clash_db, +): + db, index, vecs = clash_db + for q in (0, 5, 40, 50): + for record, score in index.find_nearest(vecs[q], k=7, include_vector=False): + ref = db.lookup_by_rid(record.get_rid()) + _same_record(record, ref) + # the record's own distance property is not the search score + assert record.get("distance") == pytest.approx( + 1000.5 + record.get("id"), rel=0, abs=1e-9 + ) + assert record.get("record") == f"record-{record.get('id')}" + + +@pytest.mark.parametrize("first", ["vector", "has", "names", "dict", "other"]) +def test_first_access_does_not_change_what_the_record_holds(hit_db, first): + db, index, vecs = hit_db + record, _ = index.find_nearest(vecs[3], k=1, include_vector=False)[0] + ref = db.lookup_by_rid(record.get_rid()) + if first == "vector": + assert list(record.get("embedding")) == list(ref.get("embedding")) + elif first == "has": + assert record.has_property("embedding") + elif first == "names": + assert list(record.get_property_names()) == list(ref.get_property_names()) + elif first == "dict": + assert list(record.to_dict()) == list(ref.to_dict()) + else: + assert record.get("id") == ref.get("id") + _same_record(record, ref) + + +@BOTH +def test_a_hit_can_be_modified_and_saved(hit_db, include_vector): + db, index, vecs = hit_db + record, _ = index.find_nearest(vecs[7], k=1, include_vector=include_vector)[0] + rid = record.get_rid() + before = list(db.lookup_by_rid(rid).get("embedding")) + with db.transaction(): + record.modify().set("label", "changed").save() + after = db.lookup_by_rid(rid) + assert after.get("label") == "changed" + assert list(after.get("embedding")) == before + # the search still finds it, with the new label + again, _ = index.find_nearest(vecs[7], k=1, include_vector=include_vector)[0] + assert again.get_rid() == rid + assert again.get("label") == "changed" + + +@BOTH +def test_options_and_by_key_hits_are_full_records(hit_db, include_vector): + db, index, vecs = hit_db + wanted = [r.get_rid() for r, _ in index.find_nearest(vecs[1], k=12)][3:9] + many = [ + index.find_nearest(vecs[1], k=4, ef_search=64, include_vector=include_vector), + index.find_nearest( + vecs[1], k=4, allowed_rids=wanted, include_vector=include_vector + ), + index.find_nearest( + vecs[1], + k=4, + ef_search=64, + allowed_rids=wanted, + include_vector=include_vector, + ), + index.find_nearest_by_key(1, k=4), + ] + for hits in many: + assert len(hits) == 4 + for record, _ in hits: + _same_record(record, db.lookup_by_rid(record.get_rid())) + filtered = index.find_nearest( + vecs[1], k=4, allowed_rids=wanted, include_vector=include_vector + ) + assert {r.get_rid() for r, _ in filtered} <= set(wanted)