From 0c8a822269ceda2ab48a99538f1acd046ab68523 Mon Sep 17 00:00:00 2001
From: Aaron Paterson <apaterson@pm.me>
Date: Wed, 7 Oct 2026 09:58:11 -0500
Subject: [PATCH] Keep each MMR candidate's similarity to the selection as a
 running maximum

vec0_mmr_rerank recomputed, at every selection step, each remaining
candidate's distance to every result already selected: about k^2/2 times the
candidate count, near 285,000 distance computations for k = 50 over the 250
candidates that lambda's overfetch gives. max_sim[i] now holds candidate i's
highest similarity to the selection and is raised against each new
selection alone, so a step measures each remaining candidate once.

A maximum over the selected results is the same however it is accumulated,
and the distances are computed with the same arguments in the same order,
so the picks are unchanged: 20 queries over 5,000 cosine vectors of 768
dimensions (k from 20 to 100, lambda from 0.3 to 0.7) return identical
rowids before and after, in 468 ms against 11.4 s in total. The MMR
snapshot tests pass unchanged.
---
 sqlite-vec.c | 35 +++++++++++++++++++----------------
 1 file changed, 19 insertions(+), 16 deletions(-)

diff --git a/sqlite-vec.c b/sqlite-vec.c
index ab77083c..1bfa43f2 100644
--- a/sqlite-vec.c
+++ b/sqlite-vec.c
@@ -7761,7 +7761,7 @@ static int vec0_mmr_rerank(
     f32 *relevance = NULL;
     i64 *out_rowids = NULL;
     f32 *out_distances = NULL;
-    void **out_vectors = NULL;
+    f32 *max_sim = NULL;
     u8 *selected = NULL;
 
     // 2. Load vectors from shadow tables
@@ -7784,15 +7784,20 @@ static int vec0_mmr_rerank(
         relevance[i] = 1.0f - (topk_distances[i] / max_dist);
     }
 
-    // 4. Greedy MMR selection
+    // 4. Greedy MMR selection. max_sim[i] is candidate i's highest similarity
+    // to the results selected so far, raised against each new selection
+    // alone, so a step measures each remaining candidate once.
     out_rowids = sqlite3_malloc64(k_target * sizeof(i64));
     out_distances = sqlite3_malloc64(k_target * sizeof(f32));
-    out_vectors = sqlite3_malloc64(k_target * sizeof(void *));
+    max_sim = sqlite3_malloc64(k_used * sizeof(f32));
     selected = sqlite3_malloc64(k_used);
-    if (!out_rowids || !out_distances || !out_vectors || !selected) {
+    if (!out_rowids || !out_distances || !max_sim || !selected) {
         rc = SQLITE_NOMEM; goto cleanup;
     }
     memset(selected, 0, k_used);
+    for (i64 i = 0; i < k_used; i++) {
+        max_sim[i] = 0.0f;
+    }
 
     i64 n_selected = 0;
     for (i64 step = 0; step < k_target && step < k_used; step++) {
@@ -7802,17 +7807,8 @@ static int vec0_mmr_rerank(
         for (i64 i = 0; i < k_used; i++) {
             if (selected[i]) continue;
 
-            // max similarity to already-selected results
-            f32 max_sim = 0.0f;
-            for (i64 j = 0; j < step; j++) {
-                f32 d = vec0_compute_distance(vector_column,
-                                              vectors[i], out_vectors[j]);
-                f32 sim = 1.0f - (d / max_dist);
-                if (sim > max_sim) max_sim = sim;
-            }
-
             f32 mmr_score = mmr_lambda * relevance[i]
-                          - (1.0f - mmr_lambda) * max_sim;
+                          - (1.0f - mmr_lambda) * max_sim[i];
             if (mmr_score > best_mmr) {
                 best_mmr = mmr_score;
                 best_idx = i;
@@ -7823,8 +7819,15 @@ static int vec0_mmr_rerank(
         selected[best_idx] = 1;
         out_rowids[step] = topk_rowids[best_idx];
         out_distances[step] = topk_distances[best_idx];
-        out_vectors[step] = vectors[best_idx];
         n_selected++;
+
+        for (i64 i = 0; i < k_used; i++) {
+            if (selected[i]) continue;
+            f32 d = vec0_compute_distance(vector_column,
+                                          vectors[i], vectors[best_idx]);
+            f32 sim = 1.0f - (d / max_dist);
+            if (sim > max_sim[i]) max_sim[i] = sim;
+        }
     }
 
     // 5. Copy only the actually-selected entries back to input arrays
@@ -7844,7 +7847,7 @@ static int vec0_mmr_rerank(
     sqlite3_free(relevance);
     sqlite3_free(out_rowids);
     sqlite3_free(out_distances);
-    sqlite3_free(out_vectors);
+    sqlite3_free(max_sim);
     sqlite3_free(selected);
     return rc;
 }