simi/
batch.rs

1//! Batch processing: uses `rayon` to evaluate arrays of thousands of
2//! strings in parallel across all CPU cores.
3
4use rayon::prelude::*;
5
6use crate::algo;
7use crate::router::{resolve_intent, Algo, Intent};
8use crate::SimiError;
9
10/// A batch comparison result.
11#[derive(Clone, Debug)]
12pub struct BatchResult {
13    /// Index of the first string in the input pair.
14    pub index_a: usize,
15    /// Index of the second string in the input pair.
16    pub index_b: usize,
17    /// The computed similarity score.
18    pub score: f64,
19}
20
21/// A batch comparator that evaluates many string pairs in parallel.
22#[derive(Clone, Debug)]
23pub struct BatchComparator {
24    algorithm: Algo,
25}
26
27impl BatchComparator {
28    /// Create a new batch comparator with the specified algorithm.
29    #[inline]
30    pub fn new(algorithm: Algo) -> Self {
31        Self { algorithm }
32    }
33
34    /// Create a batch comparator pre-configured for a specific intent.
35    ///
36    /// The intent is resolved at construction time (empty strings are used
37    /// for the length heuristic). For `Intent::Auto`, the caller should
38    /// prefer `BatchComparator::auto_per_pair()` or use `new(Algo::*)`.
39    ///
40    /// ```rust
41    /// use simi::batch::BatchComparator;
42    /// use simi::router::Intent;
43    ///
44    /// let docs = vec!["doc a".to_string(), "doc b".to_string()];
45    /// let cmp = BatchComparator::for_intent(Intent::Deduplication);
46    /// let results = cmp.compare_matrix(&docs, &docs).unwrap();
47    /// ```
48    #[inline]
49    pub fn for_intent(intent: Intent) -> Self {
50        let algo = resolve_intent(intent, "", "");
51        Self { algorithm: algo }
52    }
53
54    /// Create a batch comparator that auto-detects the best algorithm.
55    ///
56    /// This uses the empty-string heuristic for construction. For
57    /// per-pair auto-detection, use `SimiFlow::compare_with_intent()`
58    /// or construct a new `BatchComparator` per pair via the length
59    /// heuristic.
60    #[inline]
61    pub fn auto() -> Self {
62        Self::for_intent(Intent::Auto)
63    }
64
65    /// Compare two arrays element-wise in parallel.
66    ///
67    /// Each pair `(A[i], B[i])` is compared using the configured algorithm.
68    /// The two slices must have the same length.
69    #[inline]
70    pub fn compare_pairs(&self, a: &[String], b: &[String]) -> Result<Vec<BatchResult>, SimiError> {
71        if a.len() != b.len() {
72            return Err(SimiError::BatchError(
73                "Input slices must have the same length".into(),
74            ));
75        }
76
77        let results: Vec<BatchResult> = a
78            .par_iter()
79            .zip(b.par_iter())
80            .enumerate()
81            .map(|(i, (sa, sb))| {
82                let score = match &self.algorithm {
83                    Algo::Levenshtein => algo::levenshtein::similarity(sa, sb),
84                    Algo::JaroWinkler => algo::jaro_winkler::similarity(sa, sb),
85                    Algo::Hamming => algo::hamming::similarity(sa, sb).unwrap_or(0.0),
86                    Algo::Jaccard(n) => algo::jaccard::similarity(sa, sb, *n),
87                    Algo::JaccardBigram => algo::jaccard::bigram_similarity(sa, sb),
88                    Algo::JaccardTrigram => algo::jaccard::trigram_similarity(sa, sb),
89                    Algo::JaccardWord => algo::jaccard::word_similarity(sa, sb),
90                    Algo::MinHash(sh, nh) => algo::minhash::compare(sa, sb, *sh, *nh),
91                    Algo::MinHashDefault => algo::minhash::compare_default(sa, sb),
92                    Algo::SimHash(sh) => algo::simhash::compare(sa, sb, *sh),
93                    Algo::SimHashDefault => algo::simhash::compare_default(sa, sb),
94                    Algo::Bm25 => algo::bm25::similarity(sa, sb),
95                    Algo::TfIdf => algo::tfidf::similarity(sa, sb),
96                };
97                BatchResult {
98                    index_a: i,
99                    index_b: i,
100                    score,
101                }
102            })
103            .collect();
104
105        Ok(results)
106    }
107
108    /// Compare one string against many candidates in parallel.
109    ///
110    /// Returns scores for `(reference, candidates[0..n])`.
111    #[inline]
112    pub fn compare_one_to_many(
113        &self,
114        reference: &str,
115        candidates: &[String],
116    ) -> Result<Vec<BatchResult>, SimiError> {
117        let results: Vec<BatchResult> = candidates
118            .par_iter()
119            .enumerate()
120            .map(|(i, candidate)| {
121                let score = match &self.algorithm {
122                    Algo::Levenshtein => algo::levenshtein::similarity(reference, candidate),
123                    Algo::JaroWinkler => algo::jaro_winkler::similarity(reference, candidate),
124                    Algo::Hamming => algo::hamming::similarity(reference, candidate).unwrap_or(0.0),
125                    Algo::Jaccard(n) => algo::jaccard::similarity(reference, candidate, *n),
126                    Algo::JaccardBigram => algo::jaccard::bigram_similarity(reference, candidate),
127                    Algo::JaccardTrigram => algo::jaccard::trigram_similarity(reference, candidate),
128                    Algo::JaccardWord => algo::jaccard::word_similarity(reference, candidate),
129                    Algo::MinHash(sh, nh) => algo::minhash::compare(reference, candidate, *sh, *nh),
130                    Algo::MinHashDefault => algo::minhash::compare_default(reference, candidate),
131                    Algo::SimHash(sh) => algo::simhash::compare(reference, candidate, *sh),
132                    Algo::SimHashDefault => algo::simhash::compare_default(reference, candidate),
133                    Algo::Bm25 => algo::bm25::similarity(reference, candidate),
134                    Algo::TfIdf => algo::tfidf::similarity(reference, candidate),
135                };
136                BatchResult {
137                    index_a: 0,
138                    index_b: i,
139                    score,
140                }
141            })
142            .collect();
143
144        Ok(results)
145    }
146
147    /// Compare all pairs in a cross-product (matrix) in parallel.
148    ///
149    /// Returns `n * m` results comparing every element in `a` against
150    /// every element in `b`.
151    #[inline]
152    pub fn compare_matrix(
153        &self,
154        a: &[String],
155        b: &[String],
156    ) -> Result<Vec<BatchResult>, SimiError> {
157        let results: Vec<BatchResult> = a
158            .par_iter()
159            .enumerate()
160            .flat_map(|(i, sa)| {
161                b.par_iter()
162                    .enumerate()
163                    .map(|(j, sb)| {
164                        let score = match &self.algorithm {
165                            Algo::Levenshtein => algo::levenshtein::similarity(sa, sb),
166                            Algo::JaroWinkler => algo::jaro_winkler::similarity(sa, sb),
167                            Algo::Hamming => algo::hamming::similarity(sa, sb).unwrap_or(0.0),
168                            Algo::Jaccard(n) => algo::jaccard::similarity(sa, sb, *n),
169                            Algo::JaccardBigram => algo::jaccard::bigram_similarity(sa, sb),
170                            Algo::JaccardTrigram => algo::jaccard::trigram_similarity(sa, sb),
171                            Algo::JaccardWord => algo::jaccard::word_similarity(sa, sb),
172                            Algo::MinHash(sh, nh) => algo::minhash::compare(sa, sb, *sh, *nh),
173                            Algo::MinHashDefault => algo::minhash::compare_default(sa, sb),
174                            Algo::SimHash(sh) => algo::simhash::compare(sa, sb, *sh),
175                            Algo::SimHashDefault => algo::simhash::compare_default(sa, sb),
176                            Algo::Bm25 => algo::bm25::similarity(sa, sb),
177                            Algo::TfIdf => algo::tfidf::similarity(sa, sb),
178                        };
179                        BatchResult {
180                            index_a: i,
181                            index_b: j,
182                            score,
183                        }
184                    })
185                    .collect::<Vec<_>>()
186            })
187            .collect();
188
189        Ok(results)
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn compare_pairs_basic() {
199        let a = vec!["hello".into(), "world".into(), "rust".into()];
200        let b = vec!["hello".into(), "word".into(), "rusty".into()];
201
202        let comparator = BatchComparator::new(Algo::Levenshtein);
203        let results = comparator.compare_pairs(&a, &b).unwrap();
204
205        assert_eq!(results.len(), 3);
206        // First pair: "hello" == "hello" → 1.0
207        assert!((results[0].score - 1.0).abs() < f64::EPSILON);
208        // Scores are normalized
209        assert!(results[0].score >= 0.0 && results[0].score <= 1.0);
210        assert!(results[1].score >= 0.0 && results[1].score <= 1.0);
211        assert!(results[2].score >= 0.0 && results[2].score <= 1.0);
212    }
213
214    #[test]
215    fn compare_one_to_many() {
216        let reference = "hello".to_string();
217        let candidates = vec!["hello".into(), "hallo".into(), "world".into()];
218
219        let comparator = BatchComparator::new(Algo::Levenshtein);
220        let results = comparator
221            .compare_one_to_many(&reference, &candidates)
222            .unwrap();
223
224        assert_eq!(results.len(), 3);
225        assert!((results[0].score - 1.0).abs() < f64::EPSILON);
226    }
227
228    #[test]
229    fn compare_matrix() {
230        let a = vec!["hello".into(), "world".into()];
231        let b = vec!["hello".into(), "word".into()];
232
233        let comparator = BatchComparator::new(Algo::Levenshtein);
234        let results = comparator.compare_matrix(&a, &b).unwrap();
235
236        // 2 x 2 = 4 results
237        assert_eq!(results.len(), 4);
238    }
239
240    #[test]
241    fn unequal_lengths_error() {
242        let a = vec!["hello".into()];
243        let b = vec!["hello".into(), "world".into()];
244
245        let comparator = BatchComparator::new(Algo::Levenshtein);
246        let result = comparator.compare_pairs(&a, &b);
247        assert!(result.is_err());
248    }
249
250    #[test]
251    fn large_batch() {
252        let size = 100;
253        let a: Vec<String> = (0..size).map(|i| format!("string {}", i)).collect();
254        let b: Vec<String> = (0..size).map(|i| format!("string {}", i + 1)).collect();
255
256        let comparator = BatchComparator::new(Algo::Levenshtein);
257        let results = comparator.compare_pairs(&a, &b).unwrap();
258
259        assert_eq!(results.len(), size);
260        // All scores should be valid
261        for r in &results {
262            assert!(r.score >= 0.0 && r.score <= 1.0);
263        }
264    }
265}
simi/batch.rs

simi/
batch.rs
