Native find_nearest: send the query vector as one float32 array, not a Python list

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:

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

tests/test_find_nearest_query_forms.py

@@ -0,0 +1,102 @@
+"""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 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)

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) == []