diff --git a/qdrant_client/local/local_collection.py b/qdrant_client/local/local_collection.py index 6862805f..aaf5c8dd 100644 --- a/qdrant_client/local/local_collection.py +++ b/qdrant_client/local/local_collection.py @@ -67,6 +67,7 @@ from qdrant_client.local.payload_filters import ( from qdrant_client.local.payload_value_extractor import value_by_key, parse_uuid from qdrant_client.local.payload_value_setter import set_value_by_key from qdrant_client.local.persistence import CollectionPersistence +from qdrant_client.local.utils import last_argmax, swap_remove from qdrant_client.local.sparse import ( empty_sparse_vector, sort_sparse_vector, @@ -2279,17 +2280,30 @@ class LocalCollection: for i in range(len(candidate_ids)): candidate_distance_matrix[(candidate_id, candidate_ids[i])] = nearest_candidates[i] - selected = [candidate_ids[0]] - pending = candidate_ids[1:] + # Core keeps the pending candidates in an insertion-ordered set and removes the chosen + # one with `swap_remove`, which moves the last element into the freed slot. It then picks + # the best candidate with `max_by_key`, which returns the *last* maximum on ties (unlike + # `np.argmax`, which returns the first). Both details are reproduced here, otherwise + # exact ties in relevance or in MMR score are resolved differently than in core. + pending = list(range(len(candidate_ids))) + + # first point is the most relevant one + seed_position = last_argmax( + [query_raw_similarities[candidate_ids[index]] for index in pending] + ) + selected = [swap_remove(pending, seed_position)] + while len(selected) < limit and len(pending) > 0: mmr_scores = [] - for pending_id in pending: - relevance_score = query_raw_similarities[pending_id] + for pending_index in pending: + relevance_score = query_raw_similarities[candidate_ids[pending_index]] similarities_to_selected = [] - for selected_id in selected: + for selected_index in selected: similarities_to_selected.append( - candidate_distance_matrix[(pending_id, selected_id)] + candidate_distance_matrix[ + (candidate_ids[pending_index], candidate_ids[selected_index]) + ] ) max_similarity_to_selected = max(similarities_to_selected) mmr_score = ( @@ -2301,10 +2315,9 @@ class LocalCollection: [np.isneginf(sim) for sim in mmr_scores] ): # no points left passing score threshold break - best_candidate_index = np.argmax(mmr_scores).item() - selected.append(pending.pop(best_candidate_index)) + selected.append(swap_remove(pending, last_argmax(mmr_scores))) - return [id_to_point[candidate_id] for candidate_id in selected] + return [id_to_point[candidate_ids[index]] for index in selected] def _rescore_with_formula( self, diff --git a/qdrant_client/local/utils.py b/qdrant_client/local/utils.py new file mode 100644 index 00000000..535b66f2 --- /dev/null +++ b/qdrant_client/local/utils.py @@ -0,0 +1,26 @@ +"""Small helpers whose semantics have to match the Rust implementation in core.""" + + +def last_argmax(values: list[float]) -> int: + """Index of the maximum value, resolving ties in favour of the *last* maximum. + + Mirrors Rust's `Iterator::max_by_key`, which core uses to pick MMR candidates. + `np.argmax` returns the first maximum instead, which orders exact ties differently. + """ + best_index = 0 + for index in range(1, len(values)): + if values[index] >= values[best_index]: + best_index = index + return best_index + + +def swap_remove(items: list[int], position: int) -> int: + """Remove and return `items[position]`, moving the last item into the freed slot. + + Mirrors `IndexSet::swap_remove`, which core uses to drop a selected MMR candidate, + and which therefore decides the order the remaining candidates are visited in. + """ + value = items[position] + items[position] = items[-1] + items.pop() + return value diff --git a/tests/congruence_tests/test_query.py b/tests/congruence_tests/test_query.py index 65f1867f..e581fa44 100644 --- a/tests/congruence_tests/test_query.py +++ b/tests/congruence_tests/test_query.py @@ -2284,6 +2284,89 @@ def test_mmr_queries(): ) +def test_mmr_tie_breaking(): + """MMR ordering when candidates tie exactly, which is where local mode used to diverge. + + Core seeds the selection with the most relevant candidate and then picks the best MMR score + with Rust's `max_by_key`, which returns the *last* maximum on ties, while `np.argmax` returns + the first one. Core also holds the pending candidates in an `IndexSet` and drops the selected + one with `swap_remove`, which moves the last candidate into the freed slot and so changes the + order candidates are visited in. + + Relevance scores are kept distinct on purpose. Core resolves a tie in *relevance* by the + order search returned the candidates in, and `Ord for ScoredPoint` compares score only, + with no tiebreak on id, so equally scored points come back in the order they sit in the + segment. `upsert_points_impl` derives that order from `AHashMap` key iteration, which is + randomly seeded per operation, so it is stable for repeated queries against one collection + but reshuffles on every rebuild of the fixture - regardless of segment count. Every test + run rebuilds the collection, so only ties in the MMR score can be asserted on. + + All coordinates are exact binary fractions, so the MMR ties are exact in f32 both locally + and in core. + """ + + def mmr_query(client: QdrantBase, query, using=None) -> models.QueryResponse: + return client.query_points( + collection_name=COLLECTION_NAME, + query=models.NearestQuery(nearest=query, mmr=models.Mmr()), + using=using, + # these points carry no payload, and grpc reports that as None while rest reports {} + with_payload=False, + limit=10, + ) + + # dense, DOT: query is [1, 0, 0, 0], so relevance is the first coordinate. + # Once id 1 is selected, ids 2 and 4 - the first and the last of the pending candidates - + # both end up with an MMR score of -0.5. `swap_remove` visits the pending candidates as + # [4, 2, 3], while an order-preserving removal would visit them as [2, 3, 4], so the two + # disagree on which of the tied candidates is the *last* maximum. + clients = init_clients( + [ + models.PointStruct(id=1, vector=[2.0, 1.0, 0.0, 0.0]), # relevance 2.0, selected first + models.PointStruct(id=2, vector=[1.0, 0.0, 0.0, 0.0]), # relevance 1.0, MMR -0.5 + models.PointStruct(id=3, vector=[0.5, 1.0, 0.0, 0.0]), # relevance 0.5, MMR -0.75 + models.PointStruct(id=4, vector=[0.25, 0.75, 0.0, 0.0]), # relevance 0.25, MMR -0.5 + ], + vectors_config=models.VectorParams(size=4, distance=models.Distance.DOT), + ) + compare_clients_results(*clients, mmr_query, query=[1.0, 0.0, 0.0, 0.0]) + + # the same tie with EUCLID, to show the tie-breaking is not specific to DOT + clients = init_clients( + [ + models.PointStruct(id=1, vector=[0.25, 0.0]), # relevance -0.0625, selected first + models.PointStruct(id=2, vector=[0.5, 0.25]), # relevance -0.3125, MMR -0.09375 + models.PointStruct(id=3, vector=[0.5, 0.5]), # relevance -0.5, MMR -0.09375 + models.PointStruct(id=4, vector=[0.75, 0.5]), # relevance -0.8125, MMR -0.15625 + ], + vectors_config=models.VectorParams(size=2, distance=models.Distance.EUCLID), + ) + compare_clients_results(*clients, mmr_query, query=[0.0, 0.0]) + + # the same tie on a MAX_SIM multivector field, where the divergence was first spotted + clients = init_clients( + [ + # the extra vector is never the closest one, it only exercises the MAX_SIM reduction + models.PointStruct( + id=1, vector={"multi": [[2.0, 1.0, 0.0, 0.0], [0.0, 0.0, -1.0, 0.0]]} + ), + models.PointStruct(id=2, vector={"multi": [[1.0, 0.0, 0.0, 0.0]]}), + models.PointStruct(id=3, vector={"multi": [[0.5, 0.5, 0.0, 0.0]]}), + models.PointStruct(id=4, vector={"multi": [[0.25, 1.0, 0.0, 0.0]]}), + ], + vectors_config={ + "multi": models.VectorParams( + size=4, + distance=models.Distance.DOT, + multivector_config=models.MultiVectorConfig( + comparator=models.MultiVectorComparator.MAX_SIM + ), + ) + }, + ) + compare_clients_results(*clients, mmr_query, query=[[1.0, 0.0, 0.0, 0.0]], using="multi") + + def test_relevance_feedback_queries(): fixture_points = generate_fixtures()