Fix KnnMerge

Author: sicheng

rust/types/src/execution/operator.rs

@@ -283,17 +283,12 @@ pub struct KnnOutput {
 }
 
 /// The `KnnMerge` operator selects the records nearest to target from the batch vectors of records
-/// which are all sorted by distance in ascending order
+/// which are all sorted by distance in ascending order. If the same record occurs multiple times
+/// only one copy will remain in the final result.
 ///
 /// # Parameters
 /// - `fetch`: The total number of records to fetch
 ///
-/// # Inputs
-/// - `batch_distances`: The batch vector of records, each sorted by distance in ascending order
-///
-/// # Outputs
-/// - `distances`: The nearest records in either vectors, sorted by distance in ascending order
-///
 /// # Usage
 /// It can be used to merge the query results from different operators
 #[derive(Clone, Debug)]
@@ -312,10 +307,16 @@ impl KnnMerge {
             .filter_map(|(idx, itr)| itr.next().map(|rec| Reverse((rec, idx))))
             .collect::<BinaryHeap<_>>();
 
-        let mut distances = Vec::new();
+        let mut distances = Vec::<RecordDistance>::with_capacity(self.fetch as usize);
         while distances.len() < self.fetch as usize {
             if let Some(Reverse((rec, idx))) = heap_dist.pop() {
-                distances.push(rec);
+                if distances.last().is_none()
+                    || distances
+                        .last()
+                        .is_some_and(|last_rec| last_rec.offset_id != rec.offset_id)
+                {
+                    distances.push(rec);
+                }
                 if let Some(next_rec) = batch_iters
                     .get_mut(idx)
                     .expect("Enumerated index should be valid")

rust/worker/src/execution/operators/knn_merge.rs

@@ -100,7 +100,7 @@ mod tests {
                 .iter()
                 .map(|record| record.offset_id)
                 .collect::<Vec<_>>(),
-            vec![1, 2, 3, 4, 7, 7, 10, 10, 12, 13]
+            vec![1, 2, 3, 4, 7, 10, 12, 13, 16, 17]
         );
     }
 }