refactor: improve RaBitQ performance

Cargo.toml

@@ -47,9 +47,7 @@ send_fd = { path = "crates/send_fd" }
 service = { path = "crates/service" }
 
 [target.'cfg(target_os = "linux")'.dependencies]
-tikv-jemallocator = { version = "0.6.0", features = [
-    "disable_initial_exec_tls",
-] }
+tikv-jemallocator.workspace = true
 
 [lints]
 workspace = true
@@ -79,6 +77,9 @@ rustix = { version = "0.38.31", features = ["fs", "mm", "net"] }
 serde = "1"
 serde_json = "1"
 thiserror = "1"
+tikv-jemallocator = { version = "0.6.0", features = [
+    "disable_initial_exec_tls",
+] }
 toml = "0.8.10"
 validator = { version = "0.18.0", features = ["derive"] }
 

crates/pyvectors/Cargo.toml

@@ -24,6 +24,9 @@ detect = { path = "../detect" }
 indexing = { path = "../indexing" }
 stoppable_rayon = { path = "../stoppable_rayon" }
 
+[target.'cfg(target_os = "linux")'.dependencies]
+tikv-jemallocator.workspace = true
+
 [lints]
 rust.unsafe_op_in_unsafe_fn = "allow"
 rust.unused_lifetimes = "warn"

crates/pyvectors/src/lib.rs

@@ -15,6 +15,10 @@ use std::path::PathBuf;
 use validator::Validate;
 use with_labels::WithLabels;
 
+#[cfg(target_os = "linux")]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 #[pymodule]
 fn vectors(m: &Bound<'_, PyModule>) -> PyResult<()> {
     detect::init();

crates/rabitq/src/quant/error.rs

@@ -15,10 +15,29 @@ impl<T, R> ErrorFlatReranker<T, R> {
     where
         R: Fn(u32) -> (Distance, T),
     {
+        let mut top = BinaryHeap::new();
+        let mut boundary = Distance::INFINITY;
+        for (Reverse(low_u), AlwaysEqual(u)) in heap.into_iter() {
+            if low_u < boundary {
+                let (dis_u, pay_u) = (rerank)(u);
+                if dis_u < boundary {
+                    top.push((dis_u, AlwaysEqual(u), AlwaysEqual(pay_u)));
+                    if top.len() > 10 {
+                        top.pop();
+                    }
+                    if top.len() == 10 {
+                        boundary = top.peek().unwrap().0;
+                    }
+                }
+            }
+        }
         Self {
             rerank,
-            heap: heap.into(),
-            cache: BinaryHeap::new(),
+            heap: BinaryHeap::new(),
+            cache: top
+                .into_iter()
+                .map(|(a, b, c)| (Reverse(a), b, c))
+                .collect(),
         }
     }
 }