find_nearest: merge the bucket hits by score, then load only the final k

src/arcadedb/vector.py

@@ -484,33 +484,42 @@ def _lookup_record_by_rid(self, rid):
             raise ArcadeDBError(f"Vector search returned missing RID: {rid}")
         return record
 
-    def _wrap_pair_results(self, pairs):
+    def _top_hits(self, pair_lists, k):
+        """The best ``k`` ``(record, score)`` over the hit lists of every bucket index.
+
+        The candidates are merged by score first and only the final ``k`` are
+        loaded and wrapped: a type with B buckets used to load and wrap B * k
+        records to keep k. The sort is stable, so equal scores keep the order of
+        the bucket indexes, as before.
+        """
         from .core import _wrap_java_record
         from .graph import _graph_calls
 
         calls = _graph_calls()
         if calls is not None:
-            # every hit's record and score in one crossing instead of
-            # getFirst, getSecond, and the record lookup each
-            hits = calls.hits(self._database._java_db, pairs)
-            records, scores = hits[0], hits[1]
+            # merge, cut to k, and load the k records and their scores in one crossing
+            hits = calls.topHits(self._database._java_db, pair_lists, k)
+            records, scores, rids = hits[0], hits[1], hits[2]
             wrapped = []
             for i, record in enumerate(records):
                 if record is None:
                     # name the hit exactly as the per-hit path does
-                    self._lookup_record_by_rid(list(pairs)[i].getFirst())
+                    self._lookup_record_by_rid(rids[i])
                 wrapped.append(
                     (_wrap_java_record(record, self._database), float(scores[i]))
                 )
             return wrapped
 
-        wrapped_results = []
-        for pair in pairs:
-            rid = pair.getFirst()
-            score = pair.getSecond()
-            wrapped_results.append((self._lookup_record_by_rid(rid), float(score)))
-
-        return wrapped_results
+        candidates = [
+            (float(pair.getSecond()), pair.getFirst())
+            for pairs in pair_lists
+            for pair in pairs
+        ]
+        # a distance sort (lower is better); stable, so ties keep the index order
+        candidates.sort(key=lambda item: item[0])
+        return [
+            (self._lookup_record_by_rid(rid), score) for score, rid in candidates[:k]
+        ]
 
     def _find_neighbor_pairs(
         self,
@@ -574,21 +583,18 @@ def _collect_search_results(
         approximate,
         ef_search=None,
     ):
-        all_results = []
-
-        for idx in self._iter_lsm_indexes():
-            pairs = self._find_neighbor_pairs(
+        pair_lists = [
+            self._find_neighbor_pairs(
                 idx,
                 java_vector=java_vector,
                 k=k,
                 allowed_rids_set=allowed_rids_set,
                 approximate=approximate,
                 ef_search=ef_search,
             )
-
-            all_results.extend(self._wrap_pair_results(pairs))
-
-        return self._sort_results(all_results)[:k]
+            for idx in self._iter_lsm_indexes()
+        ]
+        return self._top_hits(pair_lists, k)
 
     def _run_search(
         self,

src/java/com/arcadedb/python/GraphCalls.java

@@ -19,6 +19,7 @@
 import com.arcadedb.utility.Pair;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -79,4 +80,38 @@ public static Object[] hits(final Database db, final Iterable<?> pairs) {
       s[i] = scores.get(i);
     return new Object[] { records.toArray(), s };
   }
+
+  /**
+   * The best {@code k} hits across the result lists of several vector indexes (one list per bucket index): merge by score
+   * first, load only those. {@code hits} on each list costs one record lookup per hit, so a type with B buckets paid
+   * B * k lookups and kept k. Scores are distances (lower is better); the sort is stable, so ties keep the order the
+   * lists were given in, as the Python sort of the concatenated hits did. Returns
+   * {@code {Object[] records, double[] scores, Object[] rids}} for the kept hits in order; a record is null where the
+   * engine returned null, so the caller can name that hit by its RID.
+   */
+  public static Object[] topHits(final Database db, final Object[] pairLists, final int k) {
+    final List<Pair<?, ?>> all = new ArrayList<>();
+    for (final Object list : pairLists)
+      for (final Object item : (Iterable<?>) list)
+        all.add((Pair<?, ?>) item);
+    final int total = all.size();
+    final double[] score = new double[total];
+    final Integer[] order = new Integer[total];
+    for (int i = 0; i < total; i++) {
+      score[i] = ((Number) all.get(i).getSecond()).doubleValue() + 0.0; // + 0.0: -0.0 sorts equal to 0.0, as in Python
+      order[i] = i;
+    }
+    Arrays.sort(order, (a, b) -> Double.compare(score[a], score[b])); // stable for objects
+    final int n = Math.max(0, Math.min(k, total));
+    final Object[] records = new Object[n];
+    final double[] scores = new double[n];
+    final Object[] rids = new Object[n];
+    for (int i = 0; i < n; i++) {
+      final Pair<?, ?> pair = all.get(order[i]);
+      rids[i] = pair.getFirst();
+      records[i] = db.lookupByRID((RID) rids[i], true);
+      scores[i] = score[order[i]];
+    }
+    return new Object[] { records, scores, rids };
+  }
 }

tests/test_vector.py

@@ -714,7 +714,7 @@ def test_lsm_vector_search_with_filter(self, test_db):
     def test_lsm_vector_search_looks_up_all_hits_in_one_bridge_call(
         self, test_db, monkeypatch
     ):
-        """Result materialization looks every hit up in one GraphCalls.hits call."""
+        """Result materialization merges the buckets and loads the k hits in one GraphCalls.topHits call."""
 
         test_db.command("sql", "CREATE VERTEX TYPE Doc")
         test_db.command("sql", "CREATE PROPERTY Doc.name STRING")
@@ -743,9 +743,9 @@ def test_lsm_vector_search_looks_up_all_hits_in_one_bridge_call(
         hit_calls = {"count": 0}
 
         class CountingCalls:
-            def hits(self, java_db, pairs):
+            def topHits(self, java_db, pair_lists, k):
                 hit_calls["count"] += 1
-                return calls.hits(java_db, pairs)
+                return calls.topHits(java_db, pair_lists, k)
 
         monkeypatch.setattr(graph, "_GRAPH_CALLS", CountingCalls())
 

tests/test_vector_top_k_first.py

@@ -0,0 +1,109 @@
+"""``VectorIndex.find_nearest`` merges the hits of the bucket indexes by score and loads only the final k.
+
+A type with B buckets has B vector indexes. The search used to load and wrap every hit of every
+bucket (B * k record lookups) to keep k. These tests pin the result (same k records, same order, same
+scores as a stable merge of the per-bucket hit lists) and the cost (k records loaded, not B * k).
+JVM wheel only: they read the Java hit lists of the bucket indexes.
+"""
+
+import numpy as np
+import pytest
+
+BUCKETS = 4
+DIM = 16
+COUNT = 300
+
+
+@pytest.fixture
+def index(temp_db):
+    temp_db.command("sql", f"CREATE VERTEX TYPE Doc BUCKETS {BUCKETS}")
+    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(11)
+    vectors = rng.standard_normal((COUNT, DIM)).astype(np.float32)
+    with temp_db.transaction():
+        for i in range(COUNT):
+            temp_db.command(
+                "sql",
+                "INSERT INTO Doc SET id = ?, embedding = ?",
+                i,
+                vectors[i].tolist(),
+            )
+    idx = temp_db.create_vector_index(
+        "Doc", "embedding", dimensions=DIM, id_property="id"
+    )
+    assert len(list(idx._iter_lsm_indexes())) == BUCKETS
+    return idx
+
+
+def _queries(count=8):
+    rng = np.random.default_rng(5)
+    return [rng.standard_normal(DIM).astype(np.float32) for _ in range(count)]
+
+
+def _reference(idx, query, k):
+    """Stable merge of the raw per-bucket hit lists: what the search must return."""
+    from arcadedb import to_java_float_array
+
+    java_vector = to_java_float_array(query)
+    merged = []
+    for bucket_index in idx._iter_lsm_indexes():
+        pairs = idx._find_neighbor_pairs(
+            bucket_index,
+            java_vector=java_vector,
+            k=k,
+            allowed_rids_set=None,
+            approximate=False,
+        )
+        merged.extend((float(p.getSecond()), str(p.getFirst())) for p in pairs)
+    merged.sort(key=lambda item: item[0])
+    return merged[:k]
+
+
+@pytest.mark.parametrize("k", [1, 5, 10, 50, 400])
+def test_same_records_order_and_scores_as_a_merge_of_the_buckets(index, k):
+    for query in _queries():
+        hits = index.find_nearest(query, k=k)
+        expected = _reference(index, query, k)
+        assert [(float(score), record.get_rid()) for record, score in hits] == [
+            (score, rid) for score, rid in expected
+        ]
+        assert [s for _, s in hits] == sorted(s for _, s in hits)
+
+
+def test_loads_and_wraps_only_k_records(index, monkeypatch):
+    import arcadedb.core as core
+
+    wrapped = []
+    original = core._wrap_java_record
+
+    def counting(record, database):
+        wrapped.append(record)
+        return original(record, database)
+
+    monkeypatch.setattr(core, "_wrap_java_record", counting)
+    hits = index.find_nearest(_queries(1)[0], k=10)
+    assert len(hits) == 10
+    assert len(wrapped) == 10  # was BUCKETS * 10
+
+
+def test_without_the_bridge_jar_it_loads_only_k_and_answers_the_same(
+    index, monkeypatch
+):
+    import arcadedb.graph as graph
+
+    query = _queries(1)[0]
+    with_bridge = [(r.get_rid(), s) for r, s in index.find_nearest(query, k=10)]
+
+    lookups = []
+    original = type(index)._lookup_record_by_rid
+
+    def counting(self, rid):
+        lookups.append(rid)
+        return original(self, rid)
+
+    monkeypatch.setattr(graph, "_graph_calls", lambda: None)
+    monkeypatch.setattr(type(index), "_lookup_record_by_rid", counting)
+    without_bridge = [(r.get_rid(), s) for r, s in index.find_nearest(query, k=10)]
+    assert without_bridge == with_bridge
+    assert len(lookups) == 10  # was BUCKETS * 10