Native find_nearest: a hit loads its vector on first use, not with every search (draft, stacked on #523)

bindings/python/docs/api/vector.md

@@ -329,6 +329,11 @@ also the default. If you explicitly disable eager preparation, the first call to
 
 - `List[Tuple[record, float]]`: List of `(record, distance)` tuples
     - `record`: Matched ArcadeDB record object (Vertex, Document, or Edge)
+        - On the native wheel the record arrives without the indexed vector property and loads it, once, by RID, the first time the
+          vector is read (`get`, `has_property`) or the whole record is needed (`to_dict`, `get_property_names`, `modify`, `save`).
+          Values, property order and types are those of `lookup_by_rid`. A search whose hits never touch the vector does not ship
+          it, and a hit deleted since the search raises `ArcadeDBError` on that first read. A property of the record called `record`
+          or `distance` loads the same way (the search row's own columns hide them)
     - `distance`: Distance score (float), lower is better for all functions
         - Cosine: cosine distance, range [0, 2] (0 = identical)
         - Euclidean: squared Euclidean distance, range [0, ∞) (0 = identical)

bindings/python/src/arcadedb_embedded/_native/records.py

@@ -43,9 +43,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:
@@ -56,21 +64,27 @@ 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:
         return self.get(name)
 
     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:
@@ -97,6 +111,7 @@ def set(self, name: str, value: Any):
         return self
 
     def save(self):
+        self._fill()
         db = self._db
         if self._rid is None:
             row = db.command(
@@ -124,6 +139,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:
@@ -204,3 +220,15 @@ def _wrap_dict(database, record: dict, kind: str = "d"):
     kind = props.pop("@cat", kind)
     cls = {"v": Vertex, "e": Edge}.get(kind, Document)
     return cls(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

bindings/python/src/arcadedb_embedded/_native/schema_objects.py

@@ -132,6 +132,10 @@ def buildVectorGraphNow(self):
             ).close()
 
 
+def _quote_name(name: str) -> str:
+    return "`" + name.replace("`", "``") + "`"
+
+
 def _rid_text(value) -> str:
     get_rid = getattr(value, "get_rid", None)
     return get_rid() if get_rid is not None else str(value)
@@ -159,6 +163,13 @@ def _query_vector_param(query_vector):
     )
 
 
+# find_nearest: the engine's row of a hit is the record's properties plus the columns record, @rid, @type and distance
+# (ScoredRecordRow). `*` copies it, `!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 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."""
 
@@ -175,7 +186,7 @@ def _run_search(
     ):
         if approximate:
             not_yet("VectorIndex.find_nearest_approximate over PRODUCT quantization")
-        from .records import _wrap_dict
+        from .records import _wrap_partial
 
         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
@@ -184,34 +195,39 @@ def _run_search(
             options["efSearch"] = int(ef_search)
         if allowed_rids:
             options["filter"] = [_rid_text(r) for r in allowed_rids]
+        # The hits come back without the indexed vector: expand(vectorNeighbors(..)) alone sent every hit's vector twice, in the
+        # nested `record` column and again flattened into the row. _Mixin._fill loads it on first use.
+        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",
-                "SELECT expand(vectorNeighbors(?, ?, ?, ?))",
-                self._shim._name,
-                vector,
-                int(k),
-                options,
+                "sql", sql, self._shim._name, vector, int(k), options
             )
         else:
-            rows = self._database.query(
-                "sql",
-                "SELECT expand(vectorNeighbors(?, ?, ?))",
-                self._shim._name,
-                vector,
-                int(k),
-            )
+            rows = self._database.query("sql", sql, self._shim._name, vector, int(k))
+        omitted = (prop, *_LAZY_COLUMNS)
         out = []
         for row in rows:
-            nested = row.get("record")
-            record = (
-                None
-                if nested is None
-                else _wrap_dict(
-                    self._database, nested, self._database._kind_of(nested.get("@type"))
+            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(row.get("distance"))))
+            out.append((record, float(distance)))
         return self._sort_results(out)[:k]
 
     def find_nearest_approximate(self, query_vector, k=10, allowed_rids=None):

bindings/python/tests/conftest.py

@@ -29,6 +29,8 @@
     collect_ignore.append("test_native_udf.py")
     # the native library's start-up output (#477) can only be checked with the native library loaded
     collect_ignore.append("test_native_startup_runtime.py")
+    # a search hit on the native wheel fills its vector on first use (counts lookups of the native wrapper)
+    collect_ignore.append("test_native_vector_hit_lazy.py")
 if not os.path.isdir(os.path.join(os.path.dirname(os.path.dirname(__file__)), "docs")):
     # the upstream pull request branch has no docs/ directory
     collect_ignore.append("test_docs_examples.py")

bindings/python/tests/test_native_vector_hit_lazy.py

@@ -0,0 +1,78 @@
+"""Native wheel: a ``find_nearest`` hit is built without the indexed vector and fills it on first use, with one lookup.
+
+The values themselves are checked on both wheels in ``test_vector_hit_records.py``; this file counts what is shipped and looked up.
+"""
+
+import numpy as np
+import pytest
+from arcadedb_embedded.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)
+    return temp_db, index, vecs, lookups
+
+
+def test_a_search_ships_no_vector_and_looks_up_nothing(lazy_db):
+    db, index, vecs, lookups = lazy_db
+    hits = index.find_nearest(vecs[2], k=6)
+    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 == []
+
+
+def test_the_vector_is_one_lookup_per_record_however_it_is_reached(lazy_db):
+    db, index, vecs, lookups = lazy_db
+    hits = index.find_nearest(vecs[2], k=4)
+    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_vector_of_a_closed_database_is_not_looked_up(lazy_db):
+    db, index, vecs, lookups = lazy_db
+    record, _ = index.find_nearest(vecs[2], k=1)[0]
+    db.close()
+    with pytest.raises(ArcadeDBError, match="closed"):
+        record.get("embedding")
+    assert lookups == []

bindings/python/tests/test_vector_hit_records.py

@@ -0,0 +1,123 @@
+"""The records a ``find_nearest`` hit holds: the same values as a lookup by RID, on both wheels.
+
+The native wheel builds a hit without the indexed vector and fills it on first use (``test_native_vector_hit_lazy.py``
+counts that); these tests are the contract it must keep: type, property names and their order, every value, the vector
+included, and a record that can be modified and saved after a search.
+"""
+
+import numpy as np
+import pytest
+
+DIM = 8
+N_DOC = 48
+N_SUB = 16
+
+
+@pytest.fixture
+def hit_db(temp_db):
+    """Four buckets, a subtype, and properties called ``distance`` and ``record`` (the names of the search row's own columns)."""
+    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")
+    temp_db.command("sql", "CREATE PROPERTY Doc.distance DOUBLE")
+    temp_db.command("sql", "CREATE PROPERTY Doc.record STRING")
+    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 = ?, distance = ?, record = ?, embedding = ?",
+                i,
+                f"label-{i}",
+                1000.5 + i,
+                f"record-{i}",
+                vecs[i].tolist(),
+            )
+    index = temp_db.create_vector_index(
+        "Doc", "embedding", dimensions=DIM, id_property="id"
+    )
+    return temp_db, index, vecs
+
+
+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
+
+
+def test_hits_equal_a_lookup_by_rid(hit_db):
+    db, index, vecs = hit_db
+    seen_types = set()
+    for q in (0, 5, 40, 50, 60):
+        hits = index.find_nearest(vecs[q], k=7)
+        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)
+            # 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')}"
+    assert seen_types == {"Doc", "SubDoc"}
+
+
+@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)[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)
+
+
+def test_a_hit_can_be_modified_and_saved(hit_db):
+    db, index, vecs = hit_db
+    record, _ = index.find_nearest(vecs[7], k=1)[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)[0]
+    assert again.get_rid() == rid
+    assert again.get("label") == "changed"
+
+
+def test_options_and_by_key_hits_are_full_records(hit_db):
+    db, index, vecs = hit_db
+    wanted = [r.get_rid() for r, _ in index.find_nearest(vecs[1], k=12)][3:9]
+    for hits in (
+        index.find_nearest(vecs[1], k=4, ef_search=64),
+        index.find_nearest(vecs[1], k=4, allowed_rids=wanted),
+        index.find_nearest(vecs[1], k=4, ef_search=64, allowed_rids=wanted),
+        index.find_nearest_by_key(1, k=4),
+    ):
+        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)
+    assert {r.get_rid() for r, _ in filtered} <= set(wanted)