Merge pull request #129 from nnethercott/replace-minheap-for-median

Make QueryBuilder.by_item() faster with a better top k algorithm

Author: irevoire

src/reader.rs

@@ -1,4 +1,3 @@
-use std::cmp::Reverse;
 use std::collections::BinaryHeap;
 use std::iter::repeat;
 use std::marker;
@@ -388,19 +387,16 @@ impl<'t, D: Distance> Reader<'t, D> {
                 unreachable!()
             };
             let distance = D::built_distance(query_leaf, &leaf);
-            nns_distances.push(Reverse((OrderedFloat(distance), nn)));
+            nns_distances.push((OrderedFloat(distance), nn));
         }
 
-        let mut sorted_nns = BinaryHeap::from(nns_distances);
-        let capacity = opt.count.min(sorted_nns.len());
-        let mut output = Vec::with_capacity(capacity);
-        while let Some(Reverse((OrderedFloat(dist), item))) = sorted_nns.pop() {
-            if output.len() == capacity {
-                break;
-            }
+        // Get k nearest neighbors
+        let k = opt.count.min(nns_distances.len());
+        let top_k = median_based_top_k(nns_distances, k);
+        let mut output = Vec::with_capacity(top_k.len());
+        for (OrderedFloat(dist), item) in top_k {
             output.push((item, D::normalized_distance(dist, self.dimensions)));
         }
-
         Ok(output)
     }
 
@@ -606,3 +602,39 @@ pub fn item_leaf<'a, D: Distance>(
         None => Ok(None),
     }
 }
+
+// Based on https://quickwit.io/blog/top-k-complexity, implemented in https://github.com/meilisearch/arroy/pull/129
+pub fn median_based_top_k(
+    v: Vec<(OrderedFloat<f32>, u32)>,
+    k: usize,
+) -> Vec<(OrderedFloat<f32>, u32)> {
+    let mut threshold = (OrderedFloat(f32::MAX), u32::MAX);
+    let mut buffer = Vec::with_capacity(2 * k.max(1));
+
+    // prefill with no threshold checks
+    let mut v = v.into_iter();
+    buffer.extend((&mut v).take(2 * k));
+
+    for item in v {
+        if item >= threshold {
+            continue;
+        }
+        if buffer.len() == 2 * k {
+            let (_, &mut median, _) = buffer.select_nth_unstable(k - 1);
+            threshold = median;
+            buffer.truncate(k);
+        }
+
+        // avoids buffer resizing from being inlined from vec.push()
+        let uninit = buffer.spare_capacity_mut();
+        uninit[0].write(item);
+        // SAFETY: we would have panicked above already
+        unsafe {
+            buffer.set_len(buffer.len() + 1);
+        }
+    }
+
+    buffer.sort_unstable();
+    buffer.truncate(k);
+    buffer
+}

src/tests/mod.rs

@@ -1,3 +1,5 @@
+use std::cmp::Reverse;
+use std::collections::BinaryHeap;
 use std::fmt;
 
 use heed::types::LazyDecode;
@@ -101,3 +103,19 @@ fn create_database<D: Distance>() -> DatabaseHandle<D> {
 fn rng() -> StdRng {
     StdRng::from_seed(std::array::from_fn(|_| 42))
 }
+
+fn binary_heap_based_top_k<T: Ord>(v: Vec<T>, k: usize) -> Vec<T> {
+    // max to min heap
+    let v: Vec<_> = v.into_iter().map(Reverse).collect();
+
+    let mut heap = BinaryHeap::from(v);
+    let mut output = Vec::with_capacity(k);
+
+    while let Some(Reverse(item)) = heap.pop() {
+        if output.len() == k {
+            break;
+        }
+        output.push(item);
+    }
+    output
+}

src/tests/reader.rs

@@ -1,11 +1,15 @@
 use std::fmt::Display;
 use std::num::NonZeroUsize;
 
+use ordered_float::OrderedFloat;
+use proptest::collection::vec;
+use proptest::prelude::*;
 use roaring::RoaringBitmap;
 
 use super::*;
 use crate::distance::Cosine;
 use crate::distances::{Euclidean, Manhattan};
+use crate::reader::median_based_top_k;
 use crate::{ItemId, Reader, Writer};
 
 pub struct NnsRes(pub Option<Vec<(ItemId, f32)>>);
@@ -272,3 +276,21 @@ fn try_reading_in_a_non_built_database() {
     )
     "###);
 }
+
+proptest! {
+    #[test]
+    fn median_top_k_vs_binary_heap(
+        (original, k) in vec(any::<f32>(), 1..1000).prop_flat_map(|v|{
+            let k_strategy = 1..=v.len();
+            (Just(v), k_strategy)
+        })
+    ){
+        let original: Vec<(OrderedFloat<f32>, u32)> =
+            original.into_iter().enumerate().map(|(num, item)| (OrderedFloat(item), num as u32)).collect();
+
+        let u = binary_heap_based_top_k(original.clone(), k);
+        let v = median_based_top_k(original, k);
+
+        assert_eq!(u, v);
+    }
+}