(paimon-rust) branch main updated: fix(index): deduplicate vector search top-k results (#541)

[lzljs3620320]

This is an automated email from the ASF dual-hosted git repository.

JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git


The following commit(s) were added to refs/heads/main by this push:
     new 16db7a54 fix(index): deduplicate vector search top-k results (#541)
16db7a54 is described below

commit 16db7a5404b5aaf615bba509ef4526de4780bea3
Author: QuakeWang <45645138+quakew...@users.noreply.github.com>
AuthorDate: Mon Jul 20 11:07:15 2026 +0800

    fix(index): deduplicate vector search top-k results (#541)
---
 crates/paimon/src/vector_search.rs | 59 +++++++++++++++++++++++++++++++++++---
 1 file changed, 55 insertions(+), 4 deletions(-)

diff --git a/crates/paimon/src/vector_search.rs 
b/crates/paimon/src/vector_search.rs
index 2070d5e9..53205cbf 100644
--- a/crates/paimon/src/vector_search.rs
+++ b/crates/paimon/src/vector_search.rs
@@ -203,13 +203,32 @@ impl SearchResult {
         if k == 0 {
             return Self::empty();
         }
-        if self.row_ids.len() <= k {
-            return self.clone();
-        }
 
-        let mut heap = BinaryHeap::with_capacity(k + 1);
+        let mut best_by_row_id = HashMap::with_capacity(self.row_ids.len());
         for (&row_id, &score) in self.row_ids.iter().zip(&self.scores) {
             let entry = ScoredRow { row_id, score };
+            best_by_row_id
+                .entry(row_id)
+                .and_modify(|best| {
+                    if entry.is_stronger_than(best) {
+                        *best = entry;
+                    }
+                })
+                .or_insert(entry);
+        }
+
+        if best_by_row_id.len() <= k {
+            // Keep the original row order when no truncation is needed.
+            let rows = self
+                .row_ids
+                .iter()
+                .filter_map(|row_id| best_by_row_id.remove(row_id))
+                .collect();
+            return Self::from_scored_rows(rows);
+        }
+
+        let mut heap = BinaryHeap::with_capacity(k + 1);
+        for entry in best_by_row_id.into_values() {
             if heap.len() < k {
                 heap.push(entry);
             } else if heap
@@ -339,6 +358,38 @@ mod tests {
         assert!(top.row_ids.contains(&4));
     }
 
+    #[test]
+    fn test_search_result_top_k_deduplicates_overlapping_rows() {
+        let indexed = SearchResult::new(vec![1], vec![0.9]);
+        let fallback = SearchResult::new(vec![1, 2], vec![0.8, 0.7]);
+
+        let merged = indexed.or(&fallback);
+        assert_eq!(merged.row_ids, vec![1, 1, 2]);
+        assert_eq!(merged.scores, vec![0.9, 0.8, 0.7]);
+
+        let top = merged.top_k(2);
+        assert_eq!(top.row_ids, vec![1, 2]);
+        assert_eq!(top.scores, vec![0.9, 0.7]);
+    }
+
+    #[test]
+    fn test_search_result_top_k_keeps_highest_duplicate_score() {
+        let result = SearchResult::new(vec![1, 2, 1, 3], vec![0.5, 0.8, 0.9, 
0.7]);
+
+        let top = result.top_k(2);
+        assert_eq!(top.row_ids, vec![1, 2]);
+        assert_eq!(top.scores, vec![0.9, 0.8]);
+    }
+
+    #[test]
+    fn test_search_result_top_k_preserves_order_without_truncation() {
+        let result = SearchResult::new(vec![3, 1, 2], vec![0.1, 0.9, 0.5]);
+
+        let top = result.top_k(3);
+        assert_eq!(top.row_ids, result.row_ids);
+        assert_eq!(top.scores, result.scores);
+    }
+
     #[test]
     fn test_search_result_top_k_tie_breaks_by_smaller_row_id() {
         let result = SearchResult::new(vec![30, 10, 20], vec![0.9, 0.9, 0.9]);