fix(local): match core's MMR tie-breaking (#1402)

* fix(local): match core's MMR tie-breaking

Local mode ordered MMR results differently from core whenever two
candidates tied exactly, on relevance or on MMR score. The MMR score
itself was already correct; the selection rules around it were not:

* the first point was seeded from `candidate_ids[0]`, i.e. whatever
  `search` happened to return first, instead of the most relevant
  candidate;
* `np.argmax` picked the *first* maximum, while core's `max_by_key`
  returns the *last* one on ties;
* pending candidates were kept in an order-preserving list, while core
  holds them in an `IndexSet` and drops the selected one with
  `swap_remove`, which moves the last candidate into the freed slot and
  therefore changes the order candidates are visited in - and so which
  one wins a tie.

Reproduce the three rules in `_mmr`. The divergence was spotted on a
MAX_SIM multivector field with DOT, but it is specific to neither:
plain dense vectors and EUCLID diverge the same way once an exact tie
is constructed.

The added congruence tests keep relevance scores distinct on purpose:
core orders equally relevant candidates by search order, which is not
stable, so only ties in the MMR score can be asserted on.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* refactor: move utils from collection, move tests

* tests: update comments

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

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,

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

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()