diff --git a/bindings/python/docs/development/bridge.md b/bindings/python/docs/development/bridge.md index 63c41faa528..19367c4886b 100644 --- a/bindings/python/docs/development/bridge.md +++ b/bindings/python/docs/development/bridge.md @@ -13,7 +13,7 @@ sources live in `bindings/python/src/java/com/arcadedb/python/`: | `RowAccess` | Hands a row's names and values to Python in one call (`namesAndValues`), and up to N such rows per call (`nextRows`, the fallback for `to_list()` when the bridge jar lacks `TypedRows`); the values are the engine's own objects, so Python converts them with full type fidelity. Also the first row of a result set with the set closed (`firstAndClose`, behind `ResultSet.first()`) and a property or null (`propertyOrNull`, behind `Result.get()`), each in one call instead of two or three | | `TypedRows` | Up to N rows per call as one JSON string, type-faithful (`nextRows`, which also reports whether the result set is drained): DECIMAL, BigInteger, DATE, DATETIME, Instant and zoned values, and sets are tagged and restored exactly by Python, and every other value (RIDs, embedded documents, `float[]` and `double[]`, non-finite floats, maps with non-String keys, years Python cannot hold) is returned in a side array as the engine's own object and converted by Python as before. Behind `ResultSet.to_list()` (up to 1,000 rows or 4 MiB of JSON per call, whichever comes first) and `ResultSet.iter_dicts()` (one row per call, so a lazy caller reads rows when it asks) | | `FloatVectors` | Float vectors as one binary buffer instead of JSON text or a JPype call per element: `Matrix` (n vectors in one `float[]`, a row copied into each document), and `declared` (which properties of a type are `ARRAY_OF_FLOATS`); behind `insert_many()`, `insert_columns()` and `GraphBatch.create_vertices()` | -| `GraphCalls` | Every edge of `Vertex.getEdges(direction, labels)` in one call (`outEdges`, `inEdges`, `bothEdges`, behind `Vertex.get_out_edges()` and its siblings), and the record and score of every vector search hit in one call (`hits`, behind `LSMVectorIndex` search results), and the `allowed_rids` set of a vector search from RID strings in one call (`ridSet`) | +| `GraphCalls` | Every edge of `Vertex.getEdges(direction, labels)` in one call (`outEdges`, `inEdges`, `bothEdges`, behind `Vertex.get_out_edges()` and its siblings), and the record and score of every vector search hit in one call (`hits`), and the best k hits over the hit lists of all bucket indexes of a type, merged by score before any record is loaded, in one call (`topHits`, behind `LSMVectorIndex` search results: B buckets load k records, not B * k), and the `allowed_rids` set of a vector search from RID strings in one call (`ridSet`) | | `DbCalls` | One-row statements and lookups in one call: `commandNamed`, `queryNamed`, `commandPositional`, and `queryPositional` take the parameters as the varargs of a single non-overloaded method and build the map Java-side (behind `Database.command()` and `Database.query()`), and `lookupFirst` is the whole of `Database.lookup_by_key()` | | `PythonFunctions` | Registers a Python callable as an engine `SQLFunction` in the process-wide function factory (so SQL and openCypher resolve it): the arity check, one `CommandExecutionException` naming the function for any failure of the callable, and a revoked flag so an unregistered function fails clearly. Behind `Database.register_function()` and `unregister_function()`. Each call is one JNI round trip | | `PythonRecordEvents` | One listener class for the six write events (before and after create, update, delete) that forwards each to a Python callback; a `false` from a before event is the veto, and a failure is wrapped in a `DatabaseOperationException` naming the event. Behind `Database.on()` | diff --git a/bindings/python/src/arcadedb_embedded/vector.py b/bindings/python/src/arcadedb_embedded/vector.py index e5fd83734bf..f6a7f0360a2 100644 --- a/bindings/python/src/arcadedb_embedded/vector.py +++ b/bindings/python/src/arcadedb_embedded/vector.py @@ -413,33 +413,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, @@ -503,10 +512,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, @@ -514,10 +521,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/bindings/python/src/java/com/arcadedb/python/GraphCalls.java b/bindings/python/src/java/com/arcadedb/python/GraphCalls.java index 6ff683854a3..544d8edc71b 100644 --- a/bindings/python/src/java/com/arcadedb/python/GraphCalls.java +++ b/bindings/python/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/bindings/python/tests/test_vector.py b/bindings/python/tests/test_vector.py index cd4a1fe533b..be238bb0371 100644 --- a/bindings/python/tests/test_vector.py +++ b/bindings/python/tests/test_vector.py @@ -707,7 +707,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") @@ -736,9 +736,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/bindings/python/tests/test_vector_top_k_first.py b/bindings/python/tests/test_vector_top_k_first.py new file mode 100644 index 00000000000..87d20d5fe22 --- /dev/null +++ b/bindings/python/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_embedded 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_embedded.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_embedded.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