diff --git a/src/arcadedb/vector.py b/src/arcadedb/vector.py index fea4e62..b7c7186 100644 --- a/src/arcadedb/vector.py +++ b/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,10 +583,8 @@ 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, @@ -585,10 +592,9 @@ def _collect_search_results( 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, diff --git a/src/java/com/arcadedb/python/GraphCalls.java b/src/java/com/arcadedb/python/GraphCalls.java index 6ff6838..544d8ed 100644 --- a/src/java/com/arcadedb/python/GraphCalls.java +++ b/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> 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 }; + } } diff --git a/tests/test_vector.py b/tests/test_vector.py index c3cf868..3b835f9 100644 --- a/tests/test_vector.py +++ b/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()) diff --git a/tests/test_vector_top_k_first.py b/tests/test_vector_top_k_first.py new file mode 100644 index 0000000..21c738f --- /dev/null +++ b/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