perf(search): latest tweaks to boost search performance

Author: PyDataBlog

src/python/mod.rs

@@ -8,8 +8,6 @@ use pyo3::create_exception;
 use pyo3::prelude::*;
 use std::sync::Arc;
 
-// TODO: Lazily lumped all errors here but maybe more specific errors can done? like
-// FeatureExtractionError and the likes
 create_exception!(simstring_rust, SearchError, pyo3::exceptions::PyValueError);
 
 // Wrapper for FeatureExtractor trait as I can't find any direct translation.
@@ -246,7 +244,6 @@ impl PySearcher {
                 SearchError::new_err(format!("Invalid threshold: {val}"))
             }
         })?;
-        // TODO: Explore if the python bindings can handle returning references
         Ok(results.into_iter().map(|s| s.to_string()).collect())
     }
 

src/search.rs

@@ -104,20 +104,28 @@ 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 mut all_candidates: Vec<StringId> = (min_feat_size..=max_feat_size)
+        (min_feat_size..=max_feat_size)
             .into_par_iter()
-            .flat_map(|candidate_size| {
+            .map(|candidate_size| {
                 let tau =
                     self.measure
                         .minimum_common_feature_count(query_size, candidate_size, alpha);
-                self.overlap_join(query_features, tau, candidate_size)
-            })
-            .collect();
 
-        all_candidates.sort_unstable();
-        all_candidates.dedup();
+                if tau == 0 || tau > query_size {
+                    return FxHashSet::default();
+                }
 
-        all_candidates.into_iter().collect()
+                self.overlap_join(query_features, tau, candidate_size)
+                    .into_iter()
+                    .collect::<FxHashSet<StringId>>()
+            })
+            .reduce(
+                FxHashSet::default,
+                |mut acc, set| {
+                    acc.extend(set);
+                    acc
+                },
+            )
     }
 
     fn overlap_join(
@@ -136,6 +144,11 @@ impl<'db, M: Measure> Searcher<'db, M> {
             .map(|&feature| self.db.lookup_strings(candidate_size, feature))
             .collect();
 
+        let available_features = feature_sets.iter().filter(|set| set.is_some()).count();
+        if available_features < tau {
+            return Vec::new();
+        }
+
         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()));