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 7669916..cc4c0d3 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..d678b42 --- /dev/null +++ b/tests/test_find_nearest_query_forms.py @@ -0,0 +1,119 @@ +"""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 numpy as np +import pytest + +import arcadedb +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 _assert_same(got, expected, msg=""): + """Same neighbours in the same order at the same distances. + + The native build is compared exactly: that is the path these tests guard (a query that float32 would round must + answer like its float32 form, which a tolerance of float32 size could not tell apart). The JVM build can return + distances that differ in the last float32 bit between two calls with the same query values (seen in CI: + 0.9530812501907349 against 0.9530811309814453), so there distances compare to a few float32 ulps. + """ + assert [i for i, _ in got] == [i for i, _ in expected], msg + if arcadedb.__backend__ == "native": + assert got == expected, msg + else: + assert [d for _, d in got] == pytest.approx( + [d for _, d in expected], rel=1e-6, abs=1e-9 + ), msg + + +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_same(_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_same(_answer(index.find_nearest(query, k=10)), expected) + _assert_same(_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_same(_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_same(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..6b2f7c5 --- /dev/null +++ b/tests/test_native_find_nearest_params.py @@ -0,0 +1,111 @@ +"""``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 numpy as np +import pytest + +import arcadedb +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) == []