From a3521252a1213226c031bb6e6b4a3548f28c8f20 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Sun, 25 Oct 2026 17:51:06 +0900 Subject: [PATCH] find_nearest: merge the bucket hits by score, then load only the final k (B * k lookups -> k) A type with B buckets has B vector indexes; the search loaded and wrapped every hit of every bucket and kept k. GraphCalls.topHits merges the hit lists by score (stable, as the Python sort was), cuts to k, and loads those records in one crossing. Same records, order and scores; a missing record is still named, but now only when it is among the final k. New test file pins the result against a merge of the raw per-bucket lists and counts k wraps and k lookups (the no-bridge-jar fallback too). Port of humemai/arcadedb-embedded-python#534. Co-Authored-By: Claude Sonnet 5.5 --- src/arcadedb/vector.py | 48 ++++---- src/java/com/arcadedb/python/GraphCalls.java | 35 ++++++ tests/test_vector.py | 6 +- tests/test_vector_top_k_first.py | 109 +++++++++++++++++++ 4 files changed, 174 insertions(+), 24 deletions(-) create mode 100644 tests/test_vector_top_k_first.py 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