Skip to main content

thefuzz_parity/
process.rs

1//! Ranked extraction and duplicate removal.
2
3use crate::fuzz;
4use crate::utils::full_process;
5use std::marker::PhantomData;
6
7/// A scored item borrowed from the caller's choices.
8#[derive(Clone, Copy, Debug, Eq, PartialEq)]
9pub struct Match<'a> {
10    /// The unprocessed input item.
11    pub item: &'a str,
12    /// Integer similarity score.
13    pub score: u8,
14}
15
16/// Selects a scoring function for extraction and duplicate removal.
17#[derive(Clone, Copy)]
18pub enum Scorer<'s> {
19    /// Base indel similarity after default processing.
20    Ratio,
21    /// Partial indel similarity after default processing.
22    PartialRatio,
23    /// Token sort similarity.
24    TokenSortRatio,
25    /// Token set similarity.
26    TokenSetRatio,
27    /// Partial token sort similarity.
28    PartialTokenSortRatio,
29    /// Partial token set similarity.
30    PartialTokenSetRatio,
31    /// Weighted similarity.
32    WRatio,
33    /// Processed base similarity.
34    QRatio,
35    /// Weighted similarity without the U+0080 through U+00FF deletion step.
36    UwRatio,
37    /// Processed base similarity without the U+0080 through U+00FF deletion step.
38    UqRatio,
39    /// A caller-provided scorer. Inputs receive default processing first.
40    Custom(&'s dyn Fn(&str, &str) -> u8),
41}
42
43impl std::fmt::Debug for Scorer<'_> {
44    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        let name = match self {
46            Self::Ratio => "Ratio",
47            Self::PartialRatio => "PartialRatio",
48            Self::TokenSortRatio => "TokenSortRatio",
49            Self::TokenSetRatio => "TokenSetRatio",
50            Self::PartialTokenSortRatio => "PartialTokenSortRatio",
51            Self::PartialTokenSetRatio => "PartialTokenSetRatio",
52            Self::WRatio => "WRatio",
53            Self::QRatio => "QRatio",
54            Self::UwRatio => "UwRatio",
55            Self::UqRatio => "UqRatio",
56            Self::Custom(_) => "Custom",
57        };
58        formatter.write_str(name)
59    }
60}
61
62impl Scorer<'_> {
63    fn force_ascii(self) -> bool {
64        !matches!(self, Self::UwRatio | Self::UqRatio)
65    }
66
67    fn score_processed_raw(self, a: &str, b: &str) -> f64 {
68        match self {
69            Self::Ratio | Self::QRatio | Self::UqRatio => fuzz::ratio_raw(a, b) * 100.0,
70            Self::PartialRatio => fuzz::partial_ratio_raw(a, b) * 100.0,
71            Self::TokenSortRatio => fuzz::token_sort_raw(a, b, false, false, false) * 100.0,
72            Self::TokenSetRatio => fuzz::token_set_raw(a, b, false, false, false) * 100.0,
73            Self::PartialTokenSortRatio => fuzz::token_sort_raw(a, b, false, false, true) * 100.0,
74            Self::PartialTokenSetRatio => fuzz::token_set_raw(a, b, false, false, true) * 100.0,
75            Self::WRatio | Self::UwRatio => fuzz::wratio_raw(a, b, false, false) * 100.0,
76            Self::Custom(score) => f64::from(score(a, b)),
77        }
78    }
79}
80
81/// Scores choices, filters by an inclusive cutoff, and returns a stable
82/// descending ranking.
83pub fn extract<'a, I>(
84    query: &str,
85    choices: I,
86    scorer: Scorer<'_>,
87    limit: Option<usize>,
88    score_cutoff: u8,
89) -> Vec<Match<'a>>
90where
91    I: IntoIterator<Item = &'a str>,
92{
93    let query = full_process(query, scorer.force_ascii());
94    let mut matches = Vec::new();
95    for (index, item) in choices.into_iter().enumerate() {
96        let raw_score = score_choice_raw(&query, item, scorer);
97        if raw_score >= f64::from(score_cutoff) {
98            matches.push((
99                index,
100                raw_score,
101                Match {
102                    item,
103                    score: fuzz::round_score(raw_score),
104                },
105            ));
106        }
107    }
108    matches.sort_by(|left, right| {
109        right
110            .1
111            .total_cmp(&left.1)
112            .then_with(|| left.0.cmp(&right.0))
113    });
114    matches.truncate(limit.unwrap_or(usize::MAX));
115    matches.into_iter().map(|(_, _, item)| item).collect()
116}
117
118/// Returns the first highest-scoring choice at or above an inclusive cutoff.
119pub fn extract_one<'a, I>(
120    query: &str,
121    choices: I,
122    scorer: Scorer<'_>,
123    score_cutoff: u8,
124) -> Option<Match<'a>>
125where
126    I: IntoIterator<Item = &'a str>,
127{
128    let query = full_process(query, scorer.force_ascii());
129    choices
130        .into_iter()
131        .filter_map(|item| {
132            let raw_score = score_choice_raw(&query, item, scorer);
133            (raw_score >= f64::from(score_cutoff)).then_some((raw_score, item))
134        })
135        .fold(None, |best, (raw_score, item)| match best {
136            Some((best_score, _)) if best_score >= raw_score => best,
137            _ => Some((raw_score, item)),
138        })
139        .map(|(raw_score, item)| Match {
140            item,
141            score: fuzz::round_score(raw_score),
142        })
143}
144
145/// Scores choices lazily in input order and filters by an inclusive cutoff.
146pub fn extract_without_order<'a, 's, I>(
147    query: &str,
148    choices: I,
149    scorer: Scorer<'s>,
150    score_cutoff: u8,
151) -> ExtractWithoutOrder<'a, 's, I::IntoIter>
152where
153    I: IntoIterator<Item = &'a str>,
154{
155    ExtractWithoutOrder {
156        choices: choices.into_iter(),
157        query: full_process(query, scorer.force_ascii()),
158        scorer,
159        score_cutoff,
160        item_lifetime: PhantomData,
161    }
162}
163
164/// Lazy choice scores in their input order.
165pub struct ExtractWithoutOrder<'a, 's, I> {
166    choices: I,
167    query: String,
168    scorer: Scorer<'s>,
169    score_cutoff: u8,
170    item_lifetime: PhantomData<&'a str>,
171}
172
173impl<'a, I> Iterator for ExtractWithoutOrder<'a, '_, I>
174where
175    I: Iterator<Item = &'a str>,
176{
177    type Item = Match<'a>;
178
179    fn next(&mut self) -> Option<Self::Item> {
180        for item in self.choices.by_ref() {
181            let raw_score = score_choice_raw(&self.query, item, self.scorer);
182            if raw_score >= f64::from(self.score_cutoff) {
183                return Some(Match {
184                    item,
185                    score: fuzz::round_score(raw_score),
186                });
187            }
188        }
189        None
190    }
191}
192
193/// Removes similar items and returns representatives in first-occurrence order.
194#[must_use]
195pub fn dedupe(items: &[&str], threshold: u8, scorer: Scorer<'_>) -> Vec<String> {
196    let mut representatives = Vec::with_capacity(items.len());
197    for &item in items {
198        let representative = items
199            .iter()
200            .copied()
201            .filter(|&candidate| score_pair(item, candidate, scorer) >= threshold)
202            .max_by(|left, right| {
203                left.chars()
204                    .count()
205                    .cmp(&right.chars().count())
206                    .then_with(|| left.cmp(right))
207            })
208            .unwrap_or(item);
209        if !representatives.contains(&representative) {
210            representatives.push(representative);
211        }
212    }
213
214    if representatives.len() == items.len() {
215        items.iter().map(|item| (*item).to_owned()).collect()
216    } else {
217        representatives.into_iter().map(str::to_owned).collect()
218    }
219}
220
221/// Extracts up to five matches with weighted similarity and cutoff zero.
222pub fn extract_default<'a, I>(query: &str, choices: I) -> Vec<Match<'a>>
223where
224    I: IntoIterator<Item = &'a str>,
225{
226    extract(query, choices, Scorer::WRatio, Some(5), 0)
227}
228
229/// Extracts one match with weighted similarity and cutoff zero.
230pub fn extract_one_default<'a, I>(query: &str, choices: I) -> Option<Match<'a>>
231where
232    I: IntoIterator<Item = &'a str>,
233{
234    extract_one(query, choices, Scorer::WRatio, 0)
235}
236
237/// Scores choices lazily with weighted similarity and cutoff zero.
238pub fn extract_without_order_default<'a, I>(
239    query: &str,
240    choices: I,
241) -> ExtractWithoutOrder<'a, 'static, I::IntoIter>
242where
243    I: IntoIterator<Item = &'a str>,
244{
245    extract_without_order(query, choices, Scorer::WRatio, 0)
246}
247
248/// Removes duplicates at threshold 70 with token set similarity.
249#[must_use]
250pub fn dedupe_default(items: &[&str]) -> Vec<String> {
251    dedupe(items, 70, Scorer::TokenSetRatio)
252}
253
254fn score_pair(a: &str, b: &str, scorer: Scorer<'_>) -> u8 {
255    let force_ascii = scorer.force_ascii();
256    let left = full_process(a, force_ascii);
257    let right = full_process(b, force_ascii);
258    if left.is_empty() {
259        0
260    } else {
261        fuzz::round_score(scorer.score_processed_raw(&left, &right))
262    }
263}
264
265fn score_choice_raw(query: &str, item: &str, scorer: Scorer<'_>) -> f64 {
266    if query.is_empty() {
267        0.0
268    } else {
269        let processed = full_process(item, scorer.force_ascii());
270        scorer.score_processed_raw(query, &processed)
271    }
272}