Initial insert and search optimizations

Author: PyDataBlog

src/extractors/character_ngrams.rs

@@ -1,5 +1,7 @@
 use crate::FeatureExtractor;
 use lasso::{Rodeo, Spur};
+use rustc_hash::FxHashMap;
+use std::fmt::Write;
 
 #[derive(Clone)]
 pub struct CharacterNgrams {
@@ -38,26 +40,39 @@ impl FeatureExtractor for CharacterNgrams {
         }
 
         let expected_ngrams = total_len - self.n + 1;
-        let mut ngrams = Vec::with_capacity(expected_ngrams);
-
         let padding = self.endmarker.repeat(padding_len);
 
-        // collect chars once, then slice
+        // Collect chars once, then slice
         let mut all_chars = Vec::with_capacity(total_len);
         all_chars.extend(padding.chars());
         all_chars.extend(text.chars());
         all_chars.extend(padding.chars());
 
-        // Generate n-grams using efficient windowing
+        // Inline counting + interning in one pass (no intermediate Vec<String>)
+        let mut counter: FxHashMap<String, usize> = FxHashMap::default();
+        let mut result = Vec::with_capacity(expected_ngrams);
+        let mut ngram_buffer = String::with_capacity(self.n * 4);
+        let mut counted_buffer = String::with_capacity(self.n * 4 + 8);
+
         for window in all_chars.windows(self.n) {
-            // Pre-allocate string with known capacity
-            let mut ngram = String::with_capacity(self.n * 4); // Assume max 4 bytes per char
+            // Build n-gram in reusable buffer
+            ngram_buffer.clear();
             for &ch in window {
-                ngram.push(ch);
+                ngram_buffer.push(ch);
             }
-            ngrams.push(ngram);
+
+            // Count occurrence
+            let count = counter.entry(ngram_buffer.clone()).or_insert(0);
+            *count += 1;
+
+            // Build counted string and intern
+            counted_buffer.clear();
+            counted_buffer.push_str(&ngram_buffer);
+            write!(&mut counted_buffer, "{count}").unwrap();
+            result.push(interner.get_or_intern(&counted_buffer));
         }
 
-        super::append_feature_counts(interner, ngrams)
+        result.sort_unstable();
+        result
     }
 }

src/extractors/mod.rs

@@ -2,29 +2,6 @@ mod character_ngrams;
 mod word_ngrams;
 
 use lasso::{Rodeo, Spur};
-use rustc_hash::FxHashMap;
-use std::fmt::Write;
-
-/// Takes a list of features and makes each one unique by appending its occurrence count,
-/// then interns the result and returns them sorted.
-fn append_feature_counts(interner: &mut Rodeo, features: Vec<String>) -> Vec<Spur> {
-    let mut counter: FxHashMap<String, usize> = FxHashMap::default();
-    let mut unique_features = Vec::with_capacity(features.len());
-
-    for val in features {
-        let count = counter.entry(val.clone()).or_insert(0);
-        *count += 1;
-
-        let mut unique_string = String::with_capacity(val.len() + 8); // Extra space for count
-        unique_string.push_str(&val);
-        write!(&mut unique_string, "{count}",).unwrap();
-
-        unique_features.push(interner.get_or_intern(unique_string));
-    }
-
-    unique_features.sort_unstable();
-    unique_features
-}
 
 pub trait FeatureExtractor: Send + Sync {
     /// Extracts features from text, interning them and returning their IDs.

src/extractors/word_ngrams.rs

@@ -1,5 +1,7 @@
 use crate::FeatureExtractor;
 use lasso::{Rodeo, Spur};
+use rustc_hash::FxHashMap;
+use std::fmt::Write;
 
 #[derive(Clone)]
 pub struct WordNgrams {
@@ -30,25 +32,42 @@ impl FeatureExtractor for WordNgrams {
             return vec![];
         }
 
-        let tokens = text.split(&self.splitter).filter(|s| !s.is_empty());
+        let tokens: Vec<&str> = text
+            .split(&self.splitter)
+            .filter(|s| !s.is_empty())
+            .collect();
 
-        // an iterator that includes padding
-        let padded_tokens_iter = std::iter::once(self.padder.as_str())
-            .chain(tokens)
-            .chain(std::iter::once(self.padder.as_str()));
+        // Padded tokens iterator
+        let padded_tokens: Vec<&str> = std::iter::once(self.padder.as_str())
+            .chain(tokens.into_iter())
+            .chain(std::iter::once(self.padder.as_str()))
+            .collect();
 
-        // Use a buffer to collect tokens for each n-gram
-        let mut buffer: Vec<&str> = Vec::with_capacity(self.n);
-        let mut ngrams = Vec::new();
+        if padded_tokens.len() < self.n {
+            return vec![];
+        }
+
+        // Inline counting + interning in one pass
+        let mut counter: FxHashMap<String, usize> = FxHashMap::default();
+        let expected_ngrams = padded_tokens.len() - self.n + 1;
+        let mut result = Vec::with_capacity(expected_ngrams);
+        let mut counted_buffer = String::with_capacity(64);
+
+        for window in padded_tokens.windows(self.n) {
+            let ngram = window.join(" ");
+
+            // Count occurrence
+            let count = counter.entry(ngram.clone()).or_insert(0);
+            *count += 1;
 
-        for token in padded_tokens_iter {
-            buffer.push(token);
-            if buffer.len() == self.n {
-                ngrams.push(buffer.join(" "));
-                buffer.remove(0);
-            }
+            // Build counted string and intern
+            counted_buffer.clear();
+            counted_buffer.push_str(&ngram);
+            write!(&mut counted_buffer, "{count}").unwrap();
+            result.push(interner.get_or_intern(&counted_buffer));
         }
 
-        super::append_feature_counts(interner, ngrams)
+        result.sort_unstable();
+        result
     }
 }

src/search.rs

@@ -103,26 +103,46 @@ impl<'db, M: Measure> Searcher<'db, M> {
 
         let min_feat_size = self.measure.min_feature_size(query_size, alpha);
         let max_feat_size = self.measure.max_feature_size(query_size, alpha, self.db);
+        let range_size = max_feat_size.saturating_sub(min_feat_size) + 1;
 
-        (min_feat_size..=max_feat_size)
+        // For small ranges, sequential is faster than Rayon's thread pool overhead
+        if range_size <= 4 {
+            let mut all_candidates: Vec<StringId> = Vec::new();
+            for candidate_size in min_feat_size..=max_feat_size {
+                let tau =
+                    self.measure
+                        .minimum_common_feature_count(query_size, candidate_size, alpha);
+
+                if tau == 0 || tau > query_size {
+                    continue;
+                }
+
+                all_candidates.extend(self.overlap_join(query_features, tau, candidate_size));
+            }
+            // Deduplicate at the end (faster than HashSet for small collections)
+            all_candidates.sort_unstable();
+            all_candidates.dedup();
+            return all_candidates.into_iter().collect();
+        }
+
+        // For larger ranges, use parallel iteration but collect into Vec first
+        let all_candidates: Vec<StringId> = (min_feat_size..=max_feat_size)
             .into_par_iter()
-            .map(|candidate_size| {
+            .flat_map(|candidate_size| {
                 let tau =
                     self.measure
                         .minimum_common_feature_count(query_size, candidate_size, alpha);
 
                 if tau == 0 || tau > query_size {
-                    return FxHashSet::default();
+                    return Vec::new();
                 }
 
                 self.overlap_join(query_features, tau, candidate_size)
-                    .into_iter()
-                    .collect::<FxHashSet<StringId>>()
-            })
-            .reduce(FxHashSet::default, |mut acc, set| {
-                acc.extend(set);
-                acc
             })
+            .collect();
+
+        // Deduplicate - HashSet is efficient for large collections
+        all_candidates.into_iter().collect()
     }
 
     fn overlap_join(
@@ -149,7 +169,15 @@ impl<'db, M: Measure> Searcher<'db, M> {
         let mut feature_indices: Vec<usize> = (0..query_features.len()).collect();
         feature_indices.sort_unstable_by_key(|&i| feature_sets[i].map_or(usize::MAX, |s| s.len()));
 
-        let mut candidate_counts: FxHashMap<StringId, usize> = FxHashMap::default();
+        // Pre-size HashMap based on smallest feature set to reduce rehashing
+        let estimated_candidates = feature_sets
+            .iter()
+            .filter_map(|s| s.as_ref())
+            .map(|s| s.len())
+            .min()
+            .unwrap_or(16);
+        let mut candidate_counts: FxHashMap<StringId, usize> =
+            FxHashMap::with_capacity_and_hasher(estimated_candidates, Default::default());
         let mut results = Vec::new();
         let q_len = query_features.len();