new: added rayon parallelization for vector search

Author: evilsocket

Cargo.lock

@@ -39,16 +39,17 @@ simple-home-dir = "0.3.5"
 tokio = "1.38.0"
 xml-rs = "0.8.20"
 duration-string = { version = "0.4.0", optional = true }
+rayon = { version = "1.10.0", optional = true }
 glob = "0.3.1"
 
 [features]
-default = ["ollama", "groq", "openai", "fireworks"]
+default = ["ollama", "groq", "openai", "fireworks", "rayon"]
 
 ollama = ["dep:ollama-rs"]
 groq = ["dep:groq-api-rs", "dep:duration-string"]
 openai = ["dep:openai_api_rust"]
 fireworks = ["dep:openai_api_rust"]
-
+rayon = ["dep:rayon"]
 
 [profile.release]
 lto = true        # Enable link-time optimization

Cargo.toml

@@ -1,5 +1,8 @@
 use std::{collections::HashMap, time::Instant};
 
+#[cfg(feature = "rayon")]
+use rayon::prelude::*;
+
 use anyhow::Result;
 use async_trait::async_trait;
 use colored::Colorize;
@@ -89,13 +92,25 @@ impl VectorStore for NaiveVectorStore {
         println!("[{}] {} (top {})", "rag".bold(), query, top_k);
 
         let query_vector = self.embedder.embeddings(query).await?;
-        let mut distances = vec![];
         let mut results = vec![];
 
-        // TODO: parallelize?
-        for (doc_name, doc_embedding) in &self.embeddings {
-            distances.push((doc_name, metrics::cosine(&query_vector, doc_embedding)));
-        }
+        #[cfg(feature = "rayon")]
+        let mut distances: Vec<(&String, f64)> = self
+            .embeddings
+            .par_iter()
+            .map(|(doc_name, doc_embedding)| {
+                (doc_name, metrics::cosine(&query_vector, doc_embedding))
+            })
+            .collect();
+
+        #[cfg(not(feature = "rayon"))]
+        let mut distances = {
+            let mut distances = vec![];
+            for (doc_name, doc_embedding) in &self.embeddings {
+                distances.push((doc_name, metrics::cosine(&query_vector, doc_embedding)));
+            }
+            distances
+        };
 
         distances.sort_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap());